toast-api 0.1.9

//! DGM-inspired tool system for toast
//! 
//! This module provides a flexible tool system that allows the AI agent to
//! interact with the system through various tools like bash commands and file editing.

use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

pub mod bash;
pub mod editor;

/// Information about a tool that can be used by the agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolInfo {
    pub name: String,
    pub description: String,
    pub input_schema: serde_json::Value,
}

/// Base trait for all tools
#[async_trait]
pub trait Tool: Send + Sync {
    /// Get information about this tool
    fn info(&self) -> ToolInfo;
    
    /// Execute the tool with given parameters
    async fn execute(&self, params: serde_json::Value) -> Result<String>;
}

/// Tool registry that manages all available tools
pub struct ToolRegistry {
    tools: HashMap<String, Box<dyn Tool>>,
}

impl Default for ToolRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl ToolRegistry {
    /// Create a new tool registry with all built-in tools
    pub fn new() -> Self {
        let mut tools: HashMap<String, Box<dyn Tool>> = HashMap::new();
        
        // Register bash tool
        let bash_tool = bash::BashTool::new();
        tools.insert(bash_tool.info().name.clone(), Box::new(bash_tool));
        
        // Register editor tool
        let editor_tool = editor::EditorTool::new();
        tools.insert(editor_tool.info().name.clone(), Box::new(editor_tool));
        
        Self { tools }
    }
    
    /// Get a tool by name
    pub fn get(&self, name: &str) -> Option<&dyn Tool> {
        self.tools.get(name).map(|tool| tool.as_ref())
    }
    
    /// Get all available tools
    pub fn all_tools(&self) -> Vec<ToolInfo> {
        self.tools.values().map(|t| t.info()).collect()
    }
    
    /// Execute a tool by name with parameters
    pub async fn execute(&self, name: &str, params: serde_json::Value) -> Result<String> {
        match self.tools.get(name) {
            Some(tool) => tool.execute(params).await,
            None => Err(anyhow::anyhow!("Tool '{}' not found", name)),
        }
    }
}

/// Extract tool calls from a response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    pub tool: String,
    pub params: serde_json::Value,
}

/// Parse tool calls from agent response
pub fn parse_tool_calls(response: &str) -> Vec<ToolCall> {
    let mut tool_calls = Vec::new();
    
    // Look for tool invocations in the format:
    // <tool_use>
    // {"tool": "bash", "params": {"command": "ls -la"}}
    // </tool_use>
    
    let tool_use_re = regex::Regex::new(r"(?s)<tool_use>\s*(\{.*?\})\s*</tool_use>").unwrap();
    
    for cap in tool_use_re.captures_iter(response) {
        if let Ok(tool_call) = serde_json::from_str::<ToolCall>(&cap[1]) {
            tool_calls.push(tool_call);
        }
    }
    
    // Also support the existing # exec and # read_file format for backwards compatibility
    let (reads, execs) = crate::utils::extract_commands(response);
    
    // Convert read_file commands to editor tool calls
    for path in reads {
        tool_calls.push(ToolCall {
            tool: "editor".to_string(),
            params: serde_json::json!({
                "command": "view",
                "path": path
            }),
        });
    }
    
    // Convert exec commands to bash tool calls
    for cmd in execs {
        tool_calls.push(ToolCall {
            tool: "bash".to_string(),
            params: serde_json::json!({
                "command": cmd
            }),
        });
    }
    
    tool_calls
}

/// Format tool output for display
pub fn format_tool_output(tool_name: &str, output: &str) -> String {
    // Color code based on tool type
    let colored_header = match tool_name {
        "read_file" => format!("\x1b[34m[{}]\x1b[0m", tool_name.to_uppercase()),
        "bash" | "exec" => format!("\x1b[38;5;208m[{}]\x1b[0m", tool_name.to_uppercase()),
        _ => format!("[{}]", tool_name.to_uppercase()),
    };
    
    // Return formatted output without separator lines
    format!("\n{}\n{}", colored_header, output)
}