ai_agent/

agent.rs

//! Core AI-Native Code Agent implementation

use crate::config::AgentConfig;
use crate::errors::{AgentError, ToolError};
use crate::models::LanguageModel;
use crate::tools::ToolRegistry;
use crate::types::{ExecutionResult, Task, TaskComplexity, TaskPlan, TaskResult, TaskStatus};
use std::sync::Arc;
use tokio::sync::Mutex;

/// Main AI-Native Code Agent
pub struct CodeAgent {
    model: Box<dyn LanguageModel>,
    tools: Arc<Mutex<ToolRegistry>>,
    config: AgentConfig,
    _error_handler: crate::errors::ErrorHandler,
}

impl CodeAgent {
    /// Create a new agent with the given model and configuration
    pub fn new(model: Box<dyn LanguageModel>, config: AgentConfig) -> Self {
        let _error_handler = crate::errors::ErrorHandler::new(
            config.execution.max_retries,
            config.execution.retry_delay_seconds,
        );
        Self {
            model,
            tools: Arc::new(Mutex::new(ToolRegistry::new())),
            config,
            _error_handler,
        }
    }

    /// Process a task from start to finish
    pub async fn process_task(&mut self, request: &str) -> Result<TaskResult, AgentError> {
        let task_id = uuid::Uuid::new_v4().to_string();
        let task = Task {
            id: task_id.clone(),
            request: request.to_string(),
            status: TaskStatus::Pending,
            created_at: chrono::Utc::now(),
            updated_at: chrono::Utc::now(),
            result: None,
        };

        self.execute_task_internal(task).await
    }

    /// Internal task execution
    async fn execute_task_internal(&mut self, mut task: Task) -> Result<TaskResult, AgentError> {
        task.status = TaskStatus::InProgress;
        task.updated_at = chrono::Utc::now();

        // 1. Understanding phase - use the real AI model
        let plan = self.understand_task(&task.request).await?;

        tracing::info!(
            "Task plan created: {} steps estimated",
            plan.estimated_steps.unwrap_or(0)
        );

        // 2. Execution phase - use real execution
        let execution_result = self.execute_task_real(&task.id, plan.clone()).await?;

        // 3. Generate final result
        let result = TaskResult {
            success: execution_result.success,
            summary: execution_result.summary,
            details: Some(execution_result.details),
            execution_time: Some(execution_result.execution_time),
            task_plan: Some(plan),
        };

        task.result = Some(result.clone());
        task.status = if result.success {
            TaskStatus::Completed
        } else {
            TaskStatus::Failed
        };
        task.updated_at = chrono::Utc::now();

        Ok(result)
    }

    /// Register a tool with the agent
    pub async fn register_tool<T: crate::tools::Tool + 'static>(&mut self, tool: T) {
        let mut tools = self.tools.lock().await;
        tools.register(tool);
    }

    /// Get the tool registry
    pub async fn get_tools(&self) -> Arc<Mutex<ToolRegistry>> {
        self.tools.clone()
    }

    /// Get the configuration
    pub fn get_config(&self) -> &AgentConfig {
        &self.config
    }

    /// Get the model
    pub fn get_model(&self) -> &Box<dyn LanguageModel> {
        &self.model
    }

    /// Use the real understanding engine
    async fn understand_task(&self, request: &str) -> Result<TaskPlan, AgentError> {
        tracing::info!("🧠 Starting task understanding for: {}", request);

        let prompt = format!(
            "You are an intelligent coding assistant with full autonomy.

TASK TO ANALYZE: {request}

Please analyze this task and provide:
1. Your understanding of what the user wants
2. Your approach to solving it
3. Assessment of complexity (Simple/Moderate/Complex)
4. Any requirements or dependencies you identify

You have complete freedom in how to structure your response. Be thorough but concise.

Respond in this format:
UNDERSTANDING: [your understanding]
APPROACH: [your approach]
COMPLEXITY: [Simple/Moderate/Complex]
REQUIREMENTS: [any requirements or dependencies, or \"None\"]"
        );

        tracing::debug!("📝 Sending prompt to AI model");

        let response = self
            .model
            .complete(&prompt)
            .await
            .map_err(|e| AgentError::ModelError(e))?;

        tracing::debug!("🤖 AI model response: {}", response.content);

        let plan = self.parse_task_plan(&response.content)?;

        tracing::info!("📋 Task plan created - Complexity: {:?}, Steps: {}",
                      plan.complexity, plan.estimated_steps.unwrap_or(0));

        Ok(plan)
    }

    fn parse_task_plan(&self, response: &str) -> Result<TaskPlan, AgentError> {
        let mut understanding = String::new();
        let mut approach = String::new();
        let mut complexity = TaskComplexity::Moderate;
        let mut requirements = Vec::new();

        for line in response.lines() {
            let line = line.trim();
            if line.to_uppercase().starts_with("UNDERSTANDING:") {
                understanding = line[13..].trim().to_string();
            } else if line.to_uppercase().starts_with("APPROACH:") {
                approach = line[9..].trim().to_string();
            } else if line.to_uppercase().starts_with("COMPLEXITY:") {
                match line[11..].trim().to_uppercase().as_str() {
                    "SIMPLE" => complexity = TaskComplexity::Simple,
                    "COMPLEX" => complexity = TaskComplexity::Complex,
                    _ => complexity = TaskComplexity::Moderate,
                }
            } else if line.to_uppercase().starts_with("REQUIREMENTS:") {
                let req_text = line[13..].trim();
                if req_text != "None" {
                    requirements = req_text.split(',').map(|s| s.trim().to_string()).collect();
                }
            }
        }

        let estimated_steps = match complexity {
            TaskComplexity::Simple => 1,
            TaskComplexity::Moderate => 5,
            TaskComplexity::Complex => 10,
        };

        Ok(TaskPlan {
            understanding,
            approach,
            complexity,
            estimated_steps: Some(estimated_steps),
            requirements,
        })
    }

    /// Real execution using the execution engine
    async fn execute_task_real(
        &mut self,
        task_id: &str,
        plan: TaskPlan,
    ) -> Result<ExecutionResult, AgentError> {
        tracing::info!("Starting real execution for task: {}", task_id);

        // Simple execution approach
        // In a real implementation, we'd use the execution engine properly

        // For now, let's use a simple direct execution approach
        // that actually performs the task described in the plan
        self.execute_simple_task(&plan.understanding).await
    }

    /// Simple task execution based on the task understanding
    async fn execute_simple_task(
        &self,
        task_understanding: &str,
    ) -> Result<ExecutionResult, AgentError> {
        tracing::info!("Executing simple task based on understanding: {}", task_understanding);

        // Check if the task mentions file operations
        let lower_understanding = task_understanding.to_lowercase();

        if lower_understanding.contains("read") && lower_understanding.contains("file") {
            // Try to extract file path from the understanding
            if let Some(file_path) = self.extract_file_path(task_understanding) {
                match self.read_file(&file_path).await {
                    Ok(content) => {
                        return Ok(ExecutionResult {
                            success: true,
                            summary: format!("Successfully read file: {}", file_path),
                            details: content,
                            execution_time: 2,
                        });
                    }
                    Err(e) => {
                        return Ok(ExecutionResult {
                            success: false,
                            summary: format!("Failed to read file: {}", file_path),
                            details: format!("Error: {}", e),
                            execution_time: 1,
                        });
                    }
                }
            }
        }

        if lower_understanding.contains("list") && lower_understanding.contains("file") {
            // List files in current directory
            match self.list_files(".").await {
                Ok(files) => {
                    return Ok(ExecutionResult {
                        success: true,
                        summary: "Successfully listed files".to_string(),
                        details: files,
                        execution_time: 1,
                    });
                }
                Err(e) => {
                    return Ok(ExecutionResult {
                        success: false,
                        summary: "Failed to list files".to_string(),
                        details: format!("Error: {}", e),
                        execution_time: 1,
                    });
                }
            }
        }

        if lower_understanding.contains("run") && lower_understanding.contains("command") {
            // Extract and run command
            if let Some(command) = self.extract_command(task_understanding) {
                match self.run_command(&command).await {
                    Ok(output) => {
                        return Ok(ExecutionResult {
                            success: true,
                            summary: format!("Successfully ran command: {}", command),
                            details: output,
                            execution_time: 3,
                        });
                    }
                    Err(e) => {
                        return Ok(ExecutionResult {
                            success: false,
                            summary: format!("Failed to run command: {}", command),
                            details: format!("Error: {}", e),
                            execution_time: 1,
                        });
                    }
                }
            }
        }

        // Default case: just return the understanding as the result
        Ok(ExecutionResult {
            success: true,
            summary: "Task completed".to_string(),
            details: format!("AI Analysis: {}", task_understanding),
            execution_time: 1,
        })
    }

    /// Extract file path from task understanding
    fn extract_file_path(&self, text: &str) -> Option<String> {
        // Simple regex-like extraction
        let words: Vec<&str> = text.split_whitespace().collect();
        for (i, word) in words.iter().enumerate() {
            if *word == "file" && i + 1 < words.len() {
                let next_word = words[i + 1];
                if next_word.ends_with(".txt") || next_word.ends_with(".md") ||
                   next_word.ends_with(".rs") || next_word.ends_with(".toml") {
                    return Some(next_word.trim_matches('"').trim_matches('\'').to_string());
                }
            }
        }
        None
    }

    /// Extract command from task understanding
    fn extract_command(&self, text: &str) -> Option<String> {
        let lower = text.to_lowercase();
        if lower.contains("echo") {
            if let Some(start) = lower.find("echo") {
                let command_part = &text[start..];
                if let Some(end) = command_part.find(['\'', '"']) {
                    return Some(command_part[..end].trim().to_string());
                }
                return Some(command_part.trim().to_string());
            }
        }
        None
    }

    /// Read a file
    async fn read_file(&self, path: &str) -> Result<String, AgentError> {
        let content = tokio::fs::read_to_string(path)
            .await
            .map_err(|e| AgentError::ToolError(ToolError::ExecutionError(e.to_string())))?;
        Ok(content)
    }

    /// List files in directory
    async fn list_files(&self, path: &str) -> Result<String, AgentError> {
        let mut entries = tokio::fs::read_dir(path)
            .await
            .map_err(|e| AgentError::ToolError(ToolError::ExecutionError(e.to_string())))?;

        let mut files = Vec::new();
        while let Some(entry) = entries.next_entry().await
            .map_err(|e| AgentError::ToolError(ToolError::ExecutionError(e.to_string())))? {
            let name = entry.file_name().to_string_lossy().to_string();
            let metadata = entry.metadata().await
                .map_err(|e| AgentError::ToolError(ToolError::ExecutionError(e.to_string())))?;
            let file_type = if metadata.is_dir() { "DIR" } else { "FILE" };
            files.push(format!("{}: {}", file_type, name));
        }

        files.sort();
        Ok(files.join("\n"))
    }

    /// Run a command
    async fn run_command(&self, command: &str) -> Result<String, AgentError> {
        let output = tokio::process::Command::new("sh")
            .arg("-c")
            .arg(command)
            .output()
            .await
            .map_err(|e| AgentError::ToolError(ToolError::ExecutionError(e.to_string())))?;

        if output.status.success() {
            Ok(String::from_utf8_lossy(&output.stdout).to_string())
        } else {
            Ok(String::from_utf8_lossy(&output.stderr).to_string())
        }
    }
}

/// Factory function to create an agent with default tools
pub fn create_agent_with_default_tools(
    model: Box<dyn LanguageModel>,
    config: AgentConfig,
) -> CodeAgent {
    CodeAgent::new(model, config)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::MockModel;
    use crate::tools::ReadFileTool;

    #[tokio::test]
    async fn test_agent_creation() {
        let model = Box::new(MockModel::new("test".to_string()));
        let config = AgentConfig::default();
        let agent = CodeAgent::new(model, config);

        assert_eq!(agent.get_model().model_name(), "test");
    }

    #[tokio::test]
    async fn test_tool_registration() {
        let model = Box::new(MockModel::new("test".to_string()));
        let config = AgentConfig::default();
        let mut agent = CodeAgent::new(model, config);

        agent.register_tool(ReadFileTool).await;

        let tools = agent.get_tools().await;
        let tool_names = tools.lock().await.get_tool_names();
        assert!(tool_names.contains(&"read_file".to_string()));
    }
}