ceylon_next/

agent.rs

//! Agent module providing the core Agent functionality.
//!
//! This module contains the [`Agent`] struct which is the main entry point for creating
//! and running AI agents with goal-oriented capabilities, memory, and tool integration.

use crate::llm::types::Message;
use crate::llm::{LLMClient, LLMConfig, LLMResponse, ToolCall, UniversalLLMClient};
use crate::memory::{InMemoryStore, Memory, MemoryEntry};
use crate::tasks::{OutputData, TaskRequest, TaskResponse};
use crate::tools::ToolTrait;
use crate::tools::invoker::ToolInvoker;
use crate::types::AgentStatus;
use std::sync::Arc;

// Import goal structures (you'll need to add mod goal; to your lib.rs)
use crate::goal::{Goal, GoalAnalysis, GoalAnalyzer, GoalStatus};

// ============================================
// Agent Configuration
// ============================================

/// Configuration for an AI Agent.
///
/// Controls various aspects of agent behavior including retry logic,
/// timeouts, and goal analysis capabilities.
///
/// # Examples
///
/// ```rust
/// use ceylon_next::agent::AgentConfig;
///
/// // Create default configuration
/// let config = AgentConfig::default();
///
/// // Create custom configuration
/// let mut config = AgentConfig::new(5, 120);
/// config.with_goal_analysis(false);
/// ```
#[derive(Clone)]
pub struct AgentConfig {
    max_retries: u64,
    timeout: u64,
    /// Should the agent analyze goals before starting?
    analyze_goals: bool,
}

impl Default for AgentConfig {
    fn default() -> Self {
        Self {
            max_retries: 3,
            timeout: 60,
            analyze_goals: true, // Enable goal analysis by default
        }
    }
}

impl AgentConfig {
    /// Creates a new `AgentConfig` with specified retry and timeout settings.
    ///
    /// # Arguments
    ///
    /// * `max_retries` - Maximum number of retries for failed operations
    /// * `timeout` - Timeout in seconds for agent operations
    ///
    /// # Examples
    ///
    /// ```rust
    /// use ceylon_next::agent::AgentConfig;
    ///
    /// let config = AgentConfig::new(5, 120);
    /// ```
    pub fn new(max_retries: u64, timeout: u64) -> Self {
        Self {
            max_retries,
            timeout,
            analyze_goals: true,
        }
    }

    /// Enables or disables automatic goal analysis for tasks.
    ///
    /// When enabled, the agent will automatically analyze incoming tasks
    /// and break them down into structured goals with sub-goals and success criteria.
    ///
    /// # Arguments
    ///
    /// * `enabled` - Whether to enable goal analysis
    ///
    /// # Examples
    ///
    /// ```rust
    /// use ceylon_next::agent::AgentConfig;
    ///
    /// let mut config = AgentConfig::default();
    /// config.with_goal_analysis(false);
    /// ```
    pub fn with_goal_analysis(&mut self, enabled: bool) -> &mut Self {
        self.analyze_goals = enabled;
        self
    }
}

// ============================================
// Agent
// ============================================

/// An AI Agent with goal-oriented capabilities, memory, and tool integration.
///
/// `Agent` is the core component of the Ceylon framework. It manages:
/// - LLM interactions for generating responses
/// - Goal analysis and tracking
/// - Conversation history and memory
/// - Tool execution and management
///
/// # Examples
///
/// ```rust,no_run
/// use ceylon_next::agent::Agent;
/// use ceylon_next::tasks::TaskRequest;
///
/// #[tokio::main]
/// async fn main() {
///     let mut agent = Agent::new("Assistant", "openai::gpt-4");
///
///     let task = TaskRequest::new("Hello, how are you?");
///     let response = agent.run(task).await;
///
///     println!("{:?}", response.result());
/// }
/// ```
#[derive(Clone)]
pub struct Agent {
    id: String,
    name: String,
    model: String,
    system_prompt: String,
    config: AgentConfig,
    status: AgentStatus,
    llm_client: Box<UniversalLLMClient>,
    tool_invoker: ToolInvoker,
    memory: Arc<dyn Memory>,
    /// Current goal the agent is working on
    current_goal: Option<Goal>,
}

impl Agent {
    /// Creates a new `Agent` with the specified name and LLM model.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the agent (used in system prompts)
    /// * `model` - The LLM model to use (e.g., "openai::gpt-4", "claude-3-opus", "llama2")
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    ///
    /// let agent = Agent::new("MyAssistant", "openai::gpt-4");
    /// ```
    ///
    /// # Supported Models
    ///
    /// - OpenAI: "openai::gpt-4", "gpt-3.5-turbo"
    /// - Anthropic: "claude-3-opus", "claude-3-sonnet"
    /// - Ollama: "llama2", "mistral"
    ///
    /// # Panics
    ///
    /// Panics if the model format is invalid or if an API key is required but not found
    /// in environment variables. Use [`Agent::new_with_config`] to provide API keys explicitly.
    pub fn new(name: &str, model: &str) -> Self {
        let id = uuid::Uuid::new_v4().to_string();
        let llm_client = Box::new(
            UniversalLLMClient::new(model, None)
                .unwrap_or_else(|e| panic!("Failed to create LLM client: {}", e)),
        );
        let system_prompt = Self::create_default_system_prompt(name);

        Self {
            id,
            name: name.into(),
            model: model.into(),
            system_prompt,
            config: AgentConfig::default(),
            status: AgentStatus::Stopped,
            llm_client,
            tool_invoker: ToolInvoker::new(),
            memory: Arc::new(InMemoryStore::new()),
            current_goal: None,
        }
    }

    /// Creates a new `Agent` with the specified name and LLM configuration.
    ///
    /// This method allows you to provide comprehensive LLM configuration including
    /// API keys, temperature, max tokens, and other provider-specific settings.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the agent (used in system prompts)
    /// * `llm_config` - LLM configuration with model, API key, and other settings
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    /// use ceylon_next::llm::LLMConfig;
    ///
    /// // Create agent with explicit API key
    /// let config = LLMConfig::new("openai::gpt-4")
    ///     .with_api_key("your-api-key")
    ///     .with_temperature(0.7)
    ///     .with_max_tokens(2048);
    ///
    /// let agent = Agent::new_with_config("Assistant", config).unwrap();
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The model format is invalid (should be "provider::model-name")
    /// - An API key is required but not provided or found in environment variables
    /// - The LLM provider fails to initialize
    pub fn new_with_config(name: &str, llm_config: crate::llm::LLMConfig) -> Result<Self, String> {
        let id = uuid::Uuid::new_v4().to_string();
        let llm_client = Box::new(UniversalLLMClient::new_with_config(llm_config.clone())?);
        let system_prompt = Self::create_default_system_prompt(name);

        Ok(Self {
            id,
            name: name.into(),
            model: llm_config.model.clone(),
            system_prompt,
            config: AgentConfig::default(),
            status: AgentStatus::Stopped,
            llm_client,
            tool_invoker: ToolInvoker::new(),
            memory: Arc::new(InMemoryStore::new()),
            current_goal: None,
        })
    }

    /// Sets a custom memory implementation for the agent.
    ///
    /// By default, agents use [`InMemoryStore`]. You can provide a custom
    /// implementation of the [`Memory`] trait for persistent storage.
    ///
    /// # Arguments
    ///
    /// * `memory` - An Arc-wrapped implementation of the Memory trait
    ///
    /// # Examples
    ///
    /// ```rust
    /// use ceylon_next::agent::Agent;
    /// use std::sync::Arc;
    /// // Assuming you have a custom RedisMemory implementation
    /// // let redis_memory = Arc::new(RedisMemory::new());
    /// // let mut agent = Agent::new("Assistant", "openai::gpt-4");
    /// // agent.with_memory(redis_memory);
    /// ```
    pub fn with_memory(&mut self, memory: Arc<dyn Memory>) -> &mut Self {
        self.memory = memory;
        self
    }

    /// Retrieves conversation history from memory.
    ///
    /// # Arguments
    ///
    /// * `limit` - Optional limit on number of conversations to retrieve.
    ///             If `None`, returns all history.
    ///
    /// # Returns
    ///
    /// A vector of [`MemoryEntry`] objects sorted by recency (newest first).
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let agent = Agent::new("Assistant", "openai::gpt-4");
    ///
    /// // Get last 5 conversations
    /// let recent = agent.get_history(Some(5)).await;
    ///
    /// // Get all history
    /// let all = agent.get_history(None).await;
    /// # }
    /// ```
    pub async fn get_history(&self, limit: Option<usize>) -> Vec<MemoryEntry> {
        match limit {
            Some(n) => self
                .memory
                .get_recent(&self.id, n)
                .await
                .unwrap_or_default(),
            None => self
                .memory
                .get_agent_history(&self.id)
                .await
                .unwrap_or_default(),
        }
    }

    /// Searches conversation history for messages containing the query string.
    ///
    /// Performs a case-insensitive text search across all stored messages.
    ///
    /// # Arguments
    ///
    /// * `query` - The search query string
    ///
    /// # Returns
    ///
    /// A vector of [`MemoryEntry`] objects containing the query.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let agent = Agent::new("Assistant", "openai::gpt-4");
    ///
    /// let results = agent.search_memory("Python").await;
    /// println!("Found {} conversations about Python", results.len());
    /// # }
    /// ```
    pub async fn search_memory(&self, query: &str) -> Vec<MemoryEntry> {
        self.memory
            .search(&self.id, query)
            .await
            .unwrap_or_default()
    }

    /// Clears all conversation history for this agent.
    ///
    /// # Returns
    ///
    /// `Ok(())` if successful, `Err(String)` with error message if failed.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let agent = Agent::new("Assistant", "openai::gpt-4");
    ///
    /// if let Err(e) = agent.clear_memory().await {
    ///     eprintln!("Failed to clear memory: {}", e);
    /// }
    /// # }
    /// ```
    pub async fn clear_memory(&self) -> Result<(), String> {
        self.memory.clear_agent_memory(&self.id).await
    }

    /// Returns the current goal the agent is working on, if any.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    /// use ceylon_next::tasks::TaskRequest;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let mut agent = Agent::new("Assistant", "openai::gpt-4");
    ///
    /// let task = TaskRequest::new("Build a web server");
    /// agent.run(task).await;
    ///
    /// if let Some(goal) = agent.get_current_goal() {
    ///     println!("Current goal: {}", goal.description);
    /// }
    /// # }
    /// ```
    pub fn get_current_goal(&self) -> Option<&Goal> {
        self.current_goal.as_ref()
    }

    /// Manually sets a goal for the agent to work on.
    ///
    /// This is useful when you want to provide a pre-structured goal
    /// instead of having the agent analyze the task automatically.
    ///
    /// # Arguments
    ///
    /// * `goal` - The goal to set
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    /// use ceylon_next::goal::Goal;
    ///
    /// let mut agent = Agent::new("Assistant", "openai::gpt-4");
    ///
    /// let mut goal = Goal::new("Create a REST API".to_string());
    /// goal.add_sub_goal("Design endpoints".to_string(), 1);
    /// goal.add_sub_goal("Implement handlers".to_string(), 2);
    ///
    /// agent.set_goal(goal);
    /// ```
    pub fn set_goal(&mut self, goal: Goal) {
        self.current_goal = Some(goal);
    }

    fn create_default_system_prompt(agent_name: &str) -> String {
        format!(
            "You are {}, an intelligent AI assistant with goal-oriented capabilities.

Your role:
- Understand the user's ultimate goal, not just their immediate question
- Break down complex tasks into clear sub-goals
- Track progress towards completion
- Know when you've successfully completed a task

When working on tasks:
1. First understand: What is the user really trying to achieve?
2. Define success: How will we know when we're done?
3. Break it down: What are the steps to get there?
4. Execute: Work through each step systematically
5. Verify: Check that success criteria are met before finishing

When to use tools:
- Use tools when you need external information or actions
- Always use the most appropriate tool for the task
- If a tool fails, try an alternative approach
- Explain what you're doing when using tools

HOW TO USE TOOLS:
If your model supports native tool calling, use it. Otherwise, use this format:

USE_TOOL: tool_name
{{\"parameter1\": \"value1\", \"parameter2\": \"value2\"}}
---

Example:
USE_TOOL: file_saver
{{\"filename\": \"hello.py\", \"content\": \"print('Hello, World!')\"}}
---

Your personality:
- Professional but friendly
- Patient and helpful
- Honest about your limitations
- Thorough and goal-focused
- Always double-check your work

Remember: Complete the entire goal, not just part of it. Think step by step and track your progress.",
            agent_name
        )
    }

    /// Sets a custom system prompt for the agent.
    ///
    /// The system prompt defines the agent's behavior, personality, and instructions.
    /// By default, a goal-oriented system prompt is used.
    ///
    /// # Arguments
    ///
    /// * `prompt` - The system prompt text
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    ///
    /// let mut agent = Agent::new("Assistant", "openai::gpt-4");
    /// agent.with_system_prompt("You are a helpful coding assistant specializing in Rust.");
    /// ```
    pub fn with_system_prompt(&mut self, prompt: &str) -> &mut Self {
        self.system_prompt = prompt.to_string();
        self
    }

    /// Returns the current system prompt.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    ///
    /// let agent = Agent::new("Assistant", "openai::gpt-4");
    /// println!("System prompt: {}", agent.get_system_prompt());
    /// ```
    pub fn get_system_prompt(&self) -> &str {
        &self.system_prompt
    }

    /// Sets the agent configuration.
    ///
    /// # Arguments
    ///
    /// * `config` - The [`AgentConfig`] to use
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::{Agent, AgentConfig};
    ///
    /// let mut agent = Agent::new("Assistant", "openai::gpt-4");
    ///
    /// let mut config = AgentConfig::new(5, 120);
    /// config.with_goal_analysis(true);
    ///
    /// agent.with_config(config);
    /// ```
    pub fn with_config(&mut self, config: AgentConfig) -> &mut Self {
        self.config = config;
        self
    }

    /// Configures the agent with comprehensive LLM settings using LLMConfig.
    ///
    /// This allows you to set advanced LLM parameters like temperature, top_p,
    /// reasoning, provider-specific options, and more.
    ///
    /// # Arguments
    ///
    /// * `llm_config` - The [`LLMConfig`] containing comprehensive LLM settings
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    /// use ceylon_next::llm::LLMConfig;
    ///
    /// let mut agent = Agent::new("Assistant", "openai::gpt-4");
    ///
    /// // Configure with advanced settings
    /// let llm_config = LLMConfig::new("openai::gpt-4")
    ///     .with_api_key("your-api-key")
    ///     .with_temperature(0.7)
    ///     .with_max_tokens(2048)
    ///     .with_top_p(0.9);
    ///
    /// agent.with_llm_config(llm_config);
    /// ```
    ///
    /// # Provider-Specific Examples
    ///
    /// ## Azure OpenAI
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    /// use ceylon_next::llm::LLMConfig;
    ///
    /// let mut agent = Agent::new("Assistant", "azure::gpt-4");
    ///
    /// let config = LLMConfig::new("azure::gpt-4")
    ///     .with_api_key("your-azure-key")
    ///     .with_deployment_id("your-deployment-id")
    ///     .with_api_version("2024-02-01");
    ///
    /// agent.with_llm_config(config);
    /// ```
    ///
    /// ## OpenAI with Web Search
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    /// use ceylon_next::llm::LLMConfig;
    ///
    /// let mut agent = Agent::new("Assistant", "openai::gpt-4");
    ///
    /// let config = LLMConfig::new("openai::gpt-4")
    ///     .with_api_key("your-api-key")
    ///     .with_openai_web_search(true);
    ///
    /// agent.with_llm_config(config);
    /// ```
    ///
    /// ## Anthropic with Reasoning
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    /// use ceylon_next::llm::LLMConfig;
    ///
    /// let mut agent = Agent::new("Assistant", "anthropic::claude-3-opus");
    ///
    /// let config = LLMConfig::new("anthropic::claude-3-opus")
    ///     .with_api_key("your-api-key")
    ///     .with_reasoning(true)
    ///     .with_reasoning_effort("high");
    ///
    /// agent.with_llm_config(config);
    /// ```
    pub fn with_llm_config(&mut self, llm_config: LLMConfig) -> Result<&mut Self, String> {
        // Create new LLM client with the config
        let new_client = UniversalLLMClient::new_with_config(llm_config)?;

        // Update the model name if specified in config
        self.llm_client = Box::new(new_client);

        Ok(self)
    }

    /// Adds a tool to the agent's toolset.
    ///
    /// Tools extend the agent's capabilities by allowing it to perform
    /// external actions like web searches, calculations, database queries, etc.
    ///
    /// # Arguments
    ///
    /// * `tool` - Any type implementing [`ToolTrait`]
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    /// use ceylon_next::tools::ToolTrait;
    /// use serde_json::{json, Value};
    ///
    /// struct WeatherTool;
    ///
    /// impl ToolTrait for WeatherTool {
    ///     fn name(&self) -> String { "get_weather".to_string() }
    ///     fn description(&self) -> String { "Get weather for a location".to_string() }
    ///     fn input_schema(&self) -> Value { json!({"type": "object"}) }
    ///     fn execute(&self, input: Value) -> Value { json!({"temp": 72}) }
    /// }
    ///
    /// let mut agent = Agent::new("Assistant", "openai::gpt-4");
    /// agent.add_tool(WeatherTool);
    /// ```
    pub fn add_tool<T>(&mut self, tool: T) -> &mut Self
    where
        T: ToolTrait + Send + Sync + 'static,
    {
        self.tool_invoker.add_tool(tool);
        self
    }

    /// Returns a reference to the agent's tool invoker.
    ///
    /// This can be used to inspect registered tools or invoke them manually.
    pub fn get_tool_invoker(&self) -> &crate::tools::invoker::ToolInvoker {
        &self.tool_invoker
    }

    /// NEW: Analyze the task and create a goal structure
    async fn analyze_task_goal(&mut self, task: &TaskRequest) -> Result<Goal, String> {
        println!("🎯 Analyzing task to understand the goal...");

        // Create a prompt to analyze the goal
        let analysis_prompt = GoalAnalyzer::create_analysis_prompt(task.description());

        let messages = vec![
            Message {
                role: "system".to_string(),
                content: "You are a goal analysis assistant. Analyze tasks and break them into structured goals.".to_string(),
            },
            Message {
                role: "user".to_string(),
                content: analysis_prompt,
            },
        ];

        // Ask LLM to analyze
        let response = self
            .llm_client
            .complete::<LLMResponse<String>, String>(&messages, &[])
            .await
            .map_err(|e| format!("Failed to analyze goal: {}", e))?;

        // Parse the JSON response
        let analysis: GoalAnalysis = serde_json::from_str(&response.content)
            .map_err(|e| format!("Failed to parse goal analysis: {}", e))?;

        let goal = analysis.to_goal();

        println!("\n📋 Goal Analysis:");
        println!("{}", goal.get_summary());

        Ok(goal)
    }

    /// NEW: Check if current sub-goal is complete
    fn check_sub_goal_completion(&self, messages: &[Message]) -> bool {
        // Look at recent assistant messages to see if sub-goal was addressed
        // This is a simple heuristic - you can make it smarter
        let recent_content: String = messages
            .iter()
            .rev()
            .take(3)
            .filter(|m| m.role == "assistant")
            .map(|m| m.content.clone())
            .collect::<Vec<_>>()
            .join(" ");

        // Simple check: did we generate substantial content?
        !recent_content.trim().is_empty() && recent_content.len() > 50
    }

    /// NEW: Update goal progress based on conversation
    fn update_goal_progress(&mut self, messages: &[Message]) {
        // Check completion first (immutable borrow of self)
        let is_complete = self.check_sub_goal_completion(messages);

        // Now borrow goal mutably (no conflict)
        if let Some(goal) = &mut self.current_goal {
            if goal.status == GoalStatus::InProgress && is_complete {
                // Find next incomplete sub-goal
                if let Some(index) = goal.sub_goals.iter().position(|sg| !sg.completed) {
                    println!(
                        "✅ Completed sub-goal: {}",
                        goal.sub_goals[index].description
                    );
                    goal.complete_sub_goal(
                        index,
                        Some("Completed via agent interaction".to_string()),
                    );
                }
            }
        }
    }

    /// NEW: Check if we should continue working
    fn should_continue_working(&self) -> bool {
        if let Some(goal) = &self.current_goal {
            // Continue if:
            // 1. Goal is in progress
            // 2. Not all sub-goals are complete
            // 3. Not all success criteria are met
            goal.status == GoalStatus::InProgress
                && (!goal.all_sub_goals_complete() || !goal.is_successful())
        } else {
            // No goal, use default behavior
            false
        }
    }

    /// Parse text-based tool calls from LLM response.
    ///
    /// This is a fallback for models that don't support native tool calling.
    /// It looks for patterns like:
    /// USE_TOOL: tool_name
    /// {"param": "value"}
    fn parse_text_tool_calls(&self, content: &str) -> Vec<ToolCall> {
        use serde_json::Value;

        let mut tool_calls = Vec::new();

        // Look for "USE_TOOL: tool_name" pattern
        for line in content.lines() {
            if let Some(tool_name) = line.strip_prefix("USE_TOOL:").map(|s| s.trim()) {
                println!("  📝 Detected text-based tool call: {}", tool_name);

                // Look for JSON in the following lines
                let remaining = content
                    .split_once(&format!("USE_TOOL: {}", tool_name))
                    .and_then(|(_, after)| {
                        after.split_once("---").map(|(json, _)| json.to_string())
                    })
                    .or_else(|| {
                        content
                            .split_once(&format!("USE_TOOL: {}", tool_name))
                            .map(|(_, after)| {
                                // Take lines until we hit another USE_TOOL or end
                                after
                                    .lines()
                                    .take_while(|l| !l.starts_with("USE_TOOL:"))
                                    .collect::<Vec<_>>()
                                    .join("\n")
                            })
                    });

                if let Some(json_str) = remaining {
                    // Try to parse JSON from the text
                    if let Ok(input) = serde_json::from_str::<Value>(json_str.trim()) {
                        println!("  ✓ Parsed tool input: {:?}", input);
                        tool_calls.push(ToolCall {
                            name: tool_name.to_string(),
                            input,
                        });
                    } else {
                        // Try line by line
                        for potential_json in json_str.lines() {
                            let trimmed = potential_json.trim();
                            if trimmed.starts_with('{') {
                                if let Ok(input) = serde_json::from_str::<Value>(trimmed) {
                                    println!("  ✓ Parsed tool input: {:?}", input);
                                    tool_calls.push(ToolCall {
                                        name: tool_name.to_string(),
                                        input,
                                    });
                                    break;
                                }
                            }
                        }
                    }
                }
            }
        }

        tool_calls
    }

    /// Executes a task using the agent.
    ///
    /// This is the main method for running the agent. It:
    /// 1. Analyzes the task and creates a goal structure (if enabled)
    /// 2. Loads relevant conversation history from memory
    /// 3. Iteratively processes the task with the LLM
    /// 4. Invokes tools as needed
    /// 5. Tracks progress towards goal completion
    /// 6. Saves the conversation to memory
    ///
    /// # Arguments
    ///
    /// * `task` - A [`TaskRequest`] containing the user's request
    ///
    /// # Returns
    ///
    /// A [`TaskResponse`] containing the agent's final output
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::agent::Agent;
    /// use ceylon_next::tasks::{TaskRequest, OutputData};
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let mut agent = Agent::new("Assistant", "openai::gpt-4");
    ///
    ///     let task = TaskRequest::new("Write a haiku about Rust");
    ///     let response = agent.run(task).await;
    ///
    ///     match response.result() {
    ///         OutputData::Text(text) => println!("{}", text),
    ///         _ => println!("Unexpected output type"),
    ///     }
    /// }
    /// ```
    pub async fn run(&mut self, task: TaskRequest) -> TaskResponse {
        self.status = AgentStatus::Running;

        // NEW: Analyze the task and create a goal
        if self.config.analyze_goals {
            match self.analyze_task_goal(&task).await {
                Ok(mut goal) => {
                    goal.status = GoalStatus::InProgress;
                    self.current_goal = Some(goal);
                }
                Err(e) => {
                    println!(
                        "⚠️ Could not analyze goal: {}. Continuing without goal tracking.",
                        e
                    );
                }
            }
        }

        // Load recent context from memory
        let recent_memories = self.get_history(Some(3)).await;

        // Start with system prompt
        let mut messages = vec![self.create_system_message()];

        // Add context from recent conversations
        if !recent_memories.is_empty() {
            println!(
                "📚 Loading {} recent conversations from memory",
                recent_memories.len()
            );
            let context_summary = self.create_context_message(&recent_memories);
            messages.push(context_summary);
        }

        // NEW: Add goal context if we have one
        if let Some(goal) = &self.current_goal {
            messages.push(Message {
                role: "system".to_string(),
                content: format!(
                    "CURRENT GOAL:\n{}\n\nWork through each sub-goal systematically. \
                    Don't finish until all success criteria are met.",
                    goal.get_summary()
                ),
            });
        }

        // Add current task
        messages.push(task.message());

        // Get available tools
        let tool_specs = self.tool_invoker.get_tool_specs();

        // The main agent loop
        let max_iterations = 15; // Increased for goal-oriented work
        let mut iteration = 0;

        loop {
            iteration += 1;

            if iteration > max_iterations {
                println!("⚠️ Max iterations reached");
                if let Some(goal) = &mut self.current_goal {
                    goal.status = GoalStatus::Failed;
                }
                break;
            }

            println!("\n🤔 Agent thinking... (iteration {})", iteration);

            // NEW: Show current goal progress
            if let Some(goal) = &self.current_goal {
                if let Some(next_sub_goal) = goal.get_next_sub_goal() {
                    println!("📍 Working on: {}", next_sub_goal.description);
                }
            }

            let response = self
                .llm_client
                .complete::<LLMResponse<String>, String>(&messages, &tool_specs)
                .await;

            if let Err(e) = response {
                println!("❌ LLM error: {}", e);
                break;
            }

            let llm_response = response.unwrap();

            // Handle tool calls (both native and text-based)
            let mut tool_calls = llm_response.tool_calls;

            // FALLBACK: Parse text-based tool calls for models that don't support native tool calling
            if tool_calls.is_empty() {
                tool_calls = self.parse_text_tool_calls(&llm_response.content);
            }

            if !tool_calls.is_empty() {
                for tool_call in tool_calls {
                    println!("🔧 LLM wants to use tool: {}", tool_call.name);

                    let tool_result = self
                        .tool_invoker
                        .invoke(&tool_call.name, tool_call.input)
                        .await;

                    match tool_result {
                        Ok(result) => {
                            println!("  ✓ Tool succeeded: {:?}", result);

                            messages.push(Message {
                                role: "assistant".to_string(),
                                content: format!(
                                    "I used the {} tool and got this result: {}",
                                    tool_call.name,
                                    serde_json::to_string_pretty(&result).unwrap_or_default()
                                ),
                            });
                        }
                        Err(e) => {
                            println!("  ✗ Tool failed: {}", e);

                            messages.push(Message {
                                role: "assistant".to_string(),
                                content: format!("Tool '{}' failed: {}", tool_call.name, e),
                            });
                        }
                    }
                }

                // NEW: Update goal progress after tool use
                self.update_goal_progress(&messages);

                // Continue loop to let LLM think about results
                continue;
            }

            println!("💭 Agent generated response");

            messages.push(Message {
                role: "assistant".to_string(),
                content: llm_response.content.clone(),
            });

            // NEW: Update goal progress
            self.update_goal_progress(&messages);

            // NEW: Check if we should continue based on goal
            if self.should_continue_working() {
                println!("🔄 Goal not complete yet, continuing...");

                // Add a prompt to continue working
                messages.push(Message {
                    role: "user".to_string(),
                    content: "Please continue working on the remaining sub-goals.".to_string(),
                });
                continue;
            }

            // NEW: Mark goal as complete if all criteria met
            if let Some(goal) = &mut self.current_goal {
                if goal.is_successful() && goal.all_sub_goals_complete() {
                    println!("🎉 Goal completed successfully!");
                    goal.status = GoalStatus::Completed;
                } else {
                    println!("⚠️ Goal incomplete - some criteria not met");
                }
            }

            break;
        }

        self.status = AgentStatus::Stopped;

        // Save conversation to memory
        let memory_entry = MemoryEntry::new(self.id.clone(), task.id().into(), messages.clone());

        if let Err(e) = self.memory.store(memory_entry).await {
            println!("⚠️ Failed to save to memory: {}", e);
        } else {
            println!("💾 Conversation saved to memory");
        }

        let final_answer = messages
            .last()
            .map(|m| m.content.clone())
            .unwrap_or_else(|| "No response generated".to_string());

        TaskResponse::new(&task.id(), OutputData::Text(final_answer))
    }

    fn create_context_message(&self, memories: &[MemoryEntry]) -> Message {
        let mut context = String::from("PREVIOUS CONVERSATION HISTORY:\n");
        context.push_str("Use this information to answer questions about the user.\n\n");

        for (i, memory) in memories.iter().enumerate() {
            context.push_str(&format!("=== Past Conversation {} ===\n", i + 1));

            for msg in &memory.messages {
                match msg.role.as_str() {
                    "user" => {
                        context.push_str(&format!("User said: {}\n", msg.content));
                    }
                    "assistant" => {
                        if !msg.content.is_empty() && !msg.content.contains("analyze_text") {
                            context.push_str(&format!("You replied: {}\n", msg.content));
                        }
                    }
                    _ => {}
                }
            }
            context.push_str("\n");
        }

        context.push_str("Remember: Use this history to answer questions about the user.\n");

        Message {
            role: "system".to_string(),
            content: context,
        }
    }

    fn create_system_message(&self) -> Message {
        Message {
            role: "system".to_string(),
            content: self.system_prompt.clone(),
        }
    }
}