llama-cpp-v3-agent-sdk 0.1.7

use crate::error::AgentError;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

// ---------------------------------------------------------------------------
// Tool trait
// ---------------------------------------------------------------------------

/// A tool that the agent can invoke.
///
/// Implement this trait to create custom tools. Built-in tools (bash, read,
/// write, edit, glob) already ship with the crate.
pub trait Tool: Send + Sync {
    /// Unique name used by the model to invoke this tool (e.g. `"bash"`).
    fn name(&self) -> &str;

    /// One-line human-readable description shown to the model.
    fn description(&self) -> &str;

    /// JSON Schema description of the parameters object.
    /// Must be a valid JSON object schema, e.g.:
    /// ```json
    /// {
    ///     "type": "object",
    ///     "properties": {
    ///         "command": { "type": "string", "description": "Shell command" }
    ///     },
    ///     "required": ["command"]
    /// }
    /// ```
    fn parameters_schema(&self) -> serde_json::Value;

    /// Execute the tool with the given JSON arguments.
    /// Returns the tool output as a string (will be injected back into the
    /// conversation as the tool result).
    fn execute(&self, args: &serde_json::Value) -> Result<ToolResult, AgentError>;

    /// Whether this tool requires user permission before execution.
    /// Defaults to `true` for safety.
    fn requires_permission(&self) -> bool {
        true
    }

    /// Whether this tool is considered dangerous (shown as a warning).
    fn is_dangerous(&self, _args: &serde_json::Value) -> bool {
        false
    }
}

// ---------------------------------------------------------------------------
// Tool call / result types
// ---------------------------------------------------------------------------

/// A parsed tool call extracted from model output.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    /// Name of the tool to invoke.
    pub name: String,
    /// JSON arguments for the tool.
    pub arguments: serde_json::Value,
}

/// Result returned after executing a tool.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
    /// Whether the tool execution succeeded.
    pub success: bool,
    /// Output text from the tool.
    pub output: String,
}

impl ToolResult {
    pub fn ok(output: impl Into<String>) -> Self {
        Self {
            success: true,
            output: output.into(),
        }
    }

    pub fn err(output: impl Into<String>) -> Self {
        Self {
            success: false,
            output: output.into(),
        }
    }
}

// ---------------------------------------------------------------------------
// Tool registry
// ---------------------------------------------------------------------------

/// Registry of available tools, keyed by name.
pub struct ToolRegistry {
    tools: HashMap<String, Box<dyn Tool>>,
}

impl ToolRegistry {
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
        }
    }

    /// Register a tool. Overwrites any existing tool with the same name.
    pub fn register(&mut self, tool: Box<dyn Tool>) {
        let name = tool.name().to_string();
        self.tools.insert(name, tool);
    }

    /// Look up a tool by name.
    pub fn get(&self, name: &str) -> Option<&dyn Tool> {
        self.tools.get(name).map(|t| t.as_ref())
    }

    /// Execute a tool call, returning the result.
    pub fn execute(&self, call: &ToolCall) -> Result<ToolResult, AgentError> {
        let tool = self
            .get(&call.name)
            .ok_or_else(|| AgentError::ToolNotFound(call.name.clone()))?;
        tool.execute(&call.arguments)
    }

    /// Returns an iterator over all registered tools.
    pub fn iter(&self) -> impl Iterator<Item = &dyn Tool> {
        self.tools.values().map(|t| t.as_ref())
    }

    /// Generate a system prompt fragment describing all available tools.
    ///
    /// This produces JSON tool descriptions that can be appended to the system
    /// prompt so the model knows what tools are available and how to call them.
    pub fn tools_prompt(&self) -> String {
        if self.is_empty() {
            return String::new();
        }
        let mut lines = Vec::new();
        lines.push("# Tools\n".to_string());
        lines.push("You have access to the following tools. To use a tool, output a tool call in this exact format:\n".to_string());
        lines.push("<tool_call>".to_string());
        lines.push(r#"{"name": "<tool_name>", "arguments": {<json_args>}}"#.to_string());
        lines.push("</tool_call>\n".to_string());
        lines.push("Available tools:\n".to_string());

        for tool in self.iter() {
            let schema = serde_json::json!({
                "name": tool.name(),
                "description": tool.description(),
                "parameters": tool.parameters_schema(),
            });
            lines.push(format!(
                "- {}\n```json\n{}\n```\n",
                tool.name(),
                serde_json::to_string_pretty(&schema).unwrap_or_default()
            ));
        }

        lines.push("When you want to use a tool, output ONLY the <tool_call> block. You may use multiple tool calls in a single response. After each tool call, wait for the tool result before continuing.".to_string());

        lines.join("\n")
    }

    /// Total count of registered tools.
    pub fn len(&self) -> usize {
        self.tools.len()
    }

    pub fn is_empty(&self) -> bool {
        self.tools.is_empty()
    }
}

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

// ---------------------------------------------------------------------------
// Tool call parser
// ---------------------------------------------------------------------------

/// Parse `<tool_call>...</tool_call>` blocks from model output.
///
/// Returns a list of parsed tool calls and the remaining text fragments.
pub fn parse_tool_calls(text: &str) -> (Vec<ToolCall>, Vec<String>) {
    let mut calls = Vec::new();
    let mut text_parts = Vec::new();
    let mut remaining = text;

    loop {
        if let Some(start) = remaining.find("<tool_call>") {
            let before = &remaining[..start];
            if !before.trim().is_empty() {
                text_parts.push(before.trim().to_string());
            }

            let after_tag = &remaining[start + "<tool_call>".len()..];
            if let Some(end) = after_tag.find("</tool_call>") {
                let json_str = after_tag[..end].trim();
                match serde_json::from_str::<ToolCall>(json_str) {
                    Ok(call) => calls.push(call),
                    Err(e) => {
                        // Try to be lenient — maybe the model wrapped it differently
                        text_parts.push(format!("[Failed to parse tool call: {}]", e));
                    }
                }
                remaining = &after_tag[end + "</tool_call>".len()..];
            } else {
                // Unclosed tag — treat the rest as text
                text_parts.push(remaining.to_string());
                break;
            }
        } else {
            if !remaining.trim().is_empty() {
                text_parts.push(remaining.trim().to_string());
            }
            break;
        }
    }

    (calls, text_parts)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_single_tool_call() {
        let text = r#"Let me check that for you.
<tool_call>
{"name": "bash", "arguments": {"command": "ls -la"}}
</tool_call>
"#;
        let (calls, text_parts) = parse_tool_calls(text);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].name, "bash");
        assert_eq!(text_parts.len(), 1);
        assert!(text_parts[0].contains("Let me check"));
    }

    #[test]
    fn test_parse_multiple_tool_calls() {
        let text = r#"I'll read both files.
<tool_call>
{"name": "read", "arguments": {"path": "a.txt"}}
</tool_call>
And now the second one:
<tool_call>
{"name": "read", "arguments": {"path": "b.txt"}}
</tool_call>
Done."#;
        let (calls, text_parts) = parse_tool_calls(text);
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].name, "read");
        assert_eq!(calls[1].name, "read");
        assert_eq!(text_parts.len(), 3);
    }

    #[test]
    fn test_parse_no_tool_calls() {
        let text = "Just a normal response with no tools.";
        let (calls, text_parts) = parse_tool_calls(text);
        assert_eq!(calls.len(), 0);
        assert_eq!(text_parts.len(), 1);
    }
}