owlen/crates/owlen-core/src/agent.rs

//! Agentic execution loop with ReAct pattern support.
//!
//! This module provides the core agent orchestration logic that allows an LLM
//! to reason about tasks, execute tools, and observe results in an iterative loop.

use crate::mcp::{McpClient, McpToolCall, McpToolDescriptor, McpToolResponse};
use crate::provider::Provider;
use crate::types::{ChatParameters, ChatRequest, Message};
use crate::{Error, Result};
use serde::{Deserialize, Serialize};
use std::sync::Arc;

/// Maximum number of agent iterations before stopping
const DEFAULT_MAX_ITERATIONS: usize = 15;

/// Parsed response from the LLM in ReAct format
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LlmResponse {
    /// LLM wants to execute a tool
    ToolCall {
        thought: String,
        tool_name: String,
        arguments: serde_json::Value,
    },
    /// LLM has reached a final answer
    FinalAnswer { thought: String, answer: String },
    /// LLM is just reasoning without taking action
    Reasoning { thought: String },
}

/// Parse error when LLM response doesn't match expected format
#[derive(Debug, thiserror::Error)]
pub enum ParseError {
    #[error("No recognizable pattern found in response")]
    NoPattern,
    #[error("Missing required field: {0}")]
    MissingField(String),
    #[error("Invalid JSON in ACTION_INPUT: {0}")]
    InvalidJson(String),
}

/// Result of an agent execution
#[derive(Debug, Clone)]
pub struct AgentResult {
    /// Final answer from the agent
    pub answer: String,
    /// Number of iterations taken
    pub iterations: usize,
    /// All messages exchanged during execution
    pub messages: Vec<Message>,
    /// Whether the agent completed successfully
    pub success: bool,
}

/// Configuration for agent execution
#[derive(Debug, Clone)]
pub struct AgentConfig {
    /// Maximum number of iterations
    pub max_iterations: usize,
    /// Model to use for reasoning
    pub model: String,
    /// Temperature for LLM sampling
    pub temperature: Option<f32>,
    /// Max tokens per LLM call
    pub max_tokens: Option<u32>,
}

impl Default for AgentConfig {
    fn default() -> Self {
        Self {
            max_iterations: DEFAULT_MAX_ITERATIONS,
            model: "llama3.2:latest".to_string(),
            temperature: Some(0.7),
            max_tokens: Some(4096),
        }
    }
}

/// Agent executor that orchestrates the ReAct loop
pub struct AgentExecutor {
    /// LLM provider for reasoning
    llm_client: Arc<dyn Provider>,
    /// MCP client for tool execution
    tool_client: Arc<dyn McpClient>,
    /// Agent configuration
    config: AgentConfig,
}

impl AgentExecutor {
    /// Create a new agent executor
    pub fn new(
        llm_client: Arc<dyn Provider>,
        tool_client: Arc<dyn McpClient>,
        config: AgentConfig,
    ) -> Self {
        Self {
            llm_client,
            tool_client,
            config,
        }
    }

    /// Run the agent loop with the given query
    pub async fn run(&self, query: String) -> Result<AgentResult> {
        let mut messages = vec![Message::user(query)];
        let tools = self.discover_tools().await?;

        for iteration in 0..self.config.max_iterations {
            let prompt = self.build_react_prompt(&messages, &tools);
            let response = self.generate_llm_response(prompt).await?;

            match self.parse_response(&response)? {
                LlmResponse::ToolCall {
                    thought,
                    tool_name,
                    arguments,
                } => {
                    // Add assistant's reasoning
                    messages.push(Message::assistant(format!(
                        "THOUGHT: {}\nACTION: {}\nACTION_INPUT: {}",
                        thought,
                        tool_name,
                        serde_json::to_string_pretty(&arguments).unwrap_or_default()
                    )));

                    // Execute the tool
                    let result = self.execute_tool(&tool_name, arguments).await?;

                    // Add observation
                    messages.push(Message::tool(
                        tool_name.clone(),
                        format!(
                            "OBSERVATION: {}",
                            serde_json::to_string_pretty(&result.output).unwrap_or_default()
                        ),
                    ));
                }
                LlmResponse::FinalAnswer { thought, answer } => {
                    messages.push(Message::assistant(format!(
                        "THOUGHT: {}\nFINAL_ANSWER: {}",
                        thought, answer
                    )));
                    return Ok(AgentResult {
                        answer,
                        iterations: iteration + 1,
                        messages,
                        success: true,
                    });
                }
                LlmResponse::Reasoning { thought } => {
                    messages.push(Message::assistant(format!("THOUGHT: {}", thought)));
                }
            }
        }

        // Max iterations reached
        Ok(AgentResult {
            answer: "Maximum iterations reached without finding a final answer".to_string(),
            iterations: self.config.max_iterations,
            messages,
            success: false,
        })
    }

    /// Discover available tools from the MCP client
    async fn discover_tools(&self) -> Result<Vec<McpToolDescriptor>> {
        self.tool_client.list_tools().await
    }

    /// Build a ReAct-formatted prompt with available tools
    fn build_react_prompt(
        &self,
        messages: &[Message],
        tools: &[McpToolDescriptor],
    ) -> Vec<Message> {
        let mut prompt_messages = Vec::new();

        // System prompt with ReAct instructions
        let system_prompt = self.build_system_prompt(tools);
        prompt_messages.push(Message::system(system_prompt));

        // Add conversation history
        prompt_messages.extend_from_slice(messages);

        prompt_messages
    }

    /// Build the system prompt with ReAct format and tool descriptions
    fn build_system_prompt(&self, tools: &[McpToolDescriptor]) -> String {
        let mut prompt = String::from(
            "You are an AI assistant that uses the ReAct (Reasoning and Acting) pattern to solve tasks.\n\n\
            You have access to the following tools:\n\n"
        );

        for tool in tools {
            prompt.push_str(&format!("- {}: {}\n", tool.name, tool.description));
        }

        prompt.push_str(
            "\nUse the following format:\n\n\
            THOUGHT: Your reasoning about what to do next\n\
            ACTION: tool_name\n\
            ACTION_INPUT: {\"param\": \"value\"}\n\n\
            You will receive:\n\
            OBSERVATION: The result of the tool execution\n\n\
            Continue this process until you have enough information, then provide:\n\
            THOUGHT: Final reasoning\n\
            FINAL_ANSWER: Your comprehensive answer\n\n\
            Important:\n\
            - Always start with THOUGHT to explain your reasoning\n\
            - ACTION must be one of the available tools\n\
            - ACTION_INPUT must be valid JSON\n\
            - Use FINAL_ANSWER only when you have sufficient information\n",
        );

        prompt
    }

    /// Generate an LLM response
    async fn generate_llm_response(&self, messages: Vec<Message>) -> Result<String> {
        let request = ChatRequest {
            model: self.config.model.clone(),
            messages,
            parameters: ChatParameters {
                temperature: self.config.temperature,
                max_tokens: self.config.max_tokens,
                stream: false,
                ..Default::default()
            },
            tools: None,
        };

        let response = self.llm_client.chat(request).await?;
        Ok(response.message.content)
    }

    /// Parse LLM response into structured format
    pub fn parse_response(&self, text: &str) -> Result<LlmResponse> {
        let lines: Vec<&str> = text.lines().collect();
        let mut thought = String::new();
        let mut action = String::new();
        let mut action_input = String::new();
        let mut final_answer = String::new();

        let mut i = 0;
        while i < lines.len() {
            let line = lines[i].trim();

            if line.starts_with("THOUGHT:") {
                thought = line
                    .strip_prefix("THOUGHT:")
                    .unwrap_or("")
                    .trim()
                    .to_string();
                // Collect multi-line thoughts
                i += 1;
                while i < lines.len()
                    && !lines[i].trim().starts_with("ACTION")
                    && !lines[i].trim().starts_with("FINAL_ANSWER")
                {
                    if !lines[i].trim().is_empty() {
                        thought.push(' ');
                        thought.push_str(lines[i].trim());
                    }
                    i += 1;
                }
                continue;
            }

            if line.starts_with("ACTION:") {
                action = line
                    .strip_prefix("ACTION:")
                    .unwrap_or("")
                    .trim()
                    .to_string();
                i += 1;
                continue;
            }

            if line.starts_with("ACTION_INPUT:") {
                action_input = line
                    .strip_prefix("ACTION_INPUT:")
                    .unwrap_or("")
                    .trim()
                    .to_string();
                // Collect multi-line JSON
                i += 1;
                while i < lines.len()
                    && !lines[i].trim().starts_with("THOUGHT")
                    && !lines[i].trim().starts_with("ACTION")
                {
                    action_input.push(' ');
                    action_input.push_str(lines[i].trim());
                    i += 1;
                }
                continue;
            }

            if line.starts_with("FINAL_ANSWER:") {
                final_answer = line
                    .strip_prefix("FINAL_ANSWER:")
                    .unwrap_or("")
                    .trim()
                    .to_string();
                // Collect multi-line answer
                i += 1;
                while i < lines.len() {
                    if !lines[i].trim().is_empty() {
                        final_answer.push(' ');
                        final_answer.push_str(lines[i].trim());
                    }
                    i += 1;
                }
                break;
            }

            i += 1;
        }

        // Determine response type
        if !final_answer.is_empty() {
            return Ok(LlmResponse::FinalAnswer {
                thought,
                answer: final_answer,
            });
        }

        if !action.is_empty() {
            let arguments = if action_input.is_empty() {
                serde_json::json!({})
            } else {
                serde_json::from_str(&action_input)
                    .map_err(|e| Error::Agent(ParseError::InvalidJson(e.to_string()).to_string()))?
            };

            return Ok(LlmResponse::ToolCall {
                thought,
                tool_name: action,
                arguments,
            });
        }

        if !thought.is_empty() {
            return Ok(LlmResponse::Reasoning { thought });
        }

        Err(Error::Agent(ParseError::NoPattern.to_string()))
    }

    /// Execute a tool call
    async fn execute_tool(
        &self,
        tool_name: &str,
        arguments: serde_json::Value,
    ) -> Result<McpToolResponse> {
        let call = McpToolCall {
            name: tool_name.to_string(),
            arguments,
        };
        self.tool_client.call_tool(call).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mcp::test_utils::MockMcpClient;
    use crate::provider::test_utils::MockProvider;

    #[test]
    fn test_parse_tool_call() {
        let executor = AgentExecutor {
            llm_client: Arc::new(MockProvider),
            tool_client: Arc::new(MockMcpClient),
            config: AgentConfig::default(),
        };

        let text = r#"
THOUGHT: I need to search for information about Rust
ACTION: web_search
ACTION_INPUT: {"query": "Rust programming language"}
        "#;

        let result = executor.parse_response(text).unwrap();
        match result {
            LlmResponse::ToolCall {
                thought,
                tool_name,
                arguments,
            } => {
                assert!(thought.contains("search for information"));
                assert_eq!(tool_name, "web_search");
                assert_eq!(arguments["query"], "Rust programming language");
            }
            _ => panic!("Expected ToolCall"),
        }
    }

    #[test]
    fn test_parse_final_answer() {
        let executor = AgentExecutor {
            llm_client: Arc::new(MockProvider),
            tool_client: Arc::new(MockMcpClient),
            config: AgentConfig::default(),
        };

        let text = r#"
THOUGHT: I now have enough information to answer
FINAL_ANSWER: Rust is a systems programming language focused on safety and performance.
        "#;

        let result = executor.parse_response(text).unwrap();
        match result {
            LlmResponse::FinalAnswer { thought, answer } => {
                assert!(thought.contains("enough information"));
                assert!(answer.contains("Rust is a systems programming language"));
            }
            _ => panic!("Expected FinalAnswer"),
        }
    }
}