coyote-ai 0.8.3

use super::{FunctionDeclaration, JsonSchema};
use crate::config::RequestContext;

use anyhow::{Result, bail};
use indexmap::IndexMap;
use serde_json::{Value, json};

pub const TODO_FUNCTION_PREFIX: &str = "todo__";

pub fn todo_function_declarations() -> Vec<FunctionDeclaration> {
    vec![
        FunctionDeclaration {
            name: format!("{TODO_FUNCTION_PREFIX}init"),
            description: "Initialize a new todo list with a goal. Clears any existing todos."
                .to_string(),
            parameters: JsonSchema {
                type_value: Some("object".to_string()),
                properties: Some(IndexMap::from([(
                    "goal".to_string(),
                    JsonSchema {
                        type_value: Some("string".to_string()),
                        description: Some(
                            "The overall goal to achieve when all todos are completed".into(),
                        ),
                        ..Default::default()
                    },
                )])),
                required: Some(vec!["goal".to_string()]),
                ..Default::default()
            },
            agent: false,
        },
        FunctionDeclaration {
            name: format!("{TODO_FUNCTION_PREFIX}add"),
            description: "Add a new todo item to the list.".to_string(),
            parameters: JsonSchema {
                type_value: Some("object".to_string()),
                properties: Some(IndexMap::from([(
                    "task".to_string(),
                    JsonSchema {
                        type_value: Some("string".to_string()),
                        description: Some("Description of the todo task".into()),
                        ..Default::default()
                    },
                )])),
                required: Some(vec!["task".to_string()]),
                ..Default::default()
            },
            agent: false,
        },
        FunctionDeclaration {
            name: format!("{TODO_FUNCTION_PREFIX}done"),
            description: "Mark a todo item as done by its id.".to_string(),
            parameters: JsonSchema {
                type_value: Some("object".to_string()),
                properties: Some(IndexMap::from([(
                    "id".to_string(),
                    JsonSchema {
                        type_value: Some("integer".to_string()),
                        description: Some("The id of the todo item to mark as done".into()),
                        ..Default::default()
                    },
                )])),
                required: Some(vec!["id".to_string()]),
                ..Default::default()
            },
            agent: false,
        },
        FunctionDeclaration {
            name: format!("{TODO_FUNCTION_PREFIX}list"),
            description: "Display the current todo list with status of each item.".to_string(),
            parameters: JsonSchema {
                type_value: Some("object".to_string()),
                properties: Some(IndexMap::new()),
                ..Default::default()
            },
            agent: false,
        },
        FunctionDeclaration {
            name: format!("{TODO_FUNCTION_PREFIX}clear"),
            description: "Clear the entire todo list and reset the goal. Use when the current task has been canceled or invalidated.".to_string(),
            parameters: JsonSchema {
                type_value: Some("object".to_string()),
                properties: Some(IndexMap::new()),
                ..Default::default()
            },
            agent: false,
        },
    ]
}

pub fn handle_todo_tool(ctx: &mut RequestContext, cmd_name: &str, args: &Value) -> Result<Value> {
    let action = cmd_name
        .strip_prefix(TODO_FUNCTION_PREFIX)
        .unwrap_or(cmd_name);

    if !ctx.app.config.function_calling_support {
        bail!("Cannot use todo tools: function calling is disabled.");
    }
    let auto_config = ctx.auto_continue_config();
    if !auto_config.enabled {
        bail!(
            "Auto-continue is not enabled. Set 'auto_continue: true' in your config to use todo tools."
        );
    }

    match action {
        "init" => {
            let goal = args.get("goal").and_then(Value::as_str).unwrap_or_default();
            ctx.init_todo_list(goal);
            Ok(json!({"status": "ok", "message": "Initialized new todo list"}))
        }
        "add" => {
            let task = args.get("task").and_then(Value::as_str).unwrap_or_default();
            if task.is_empty() {
                return Ok(json!({"error": "task description is required"}));
            }
            let id = ctx.add_todo(task);
            Ok(json!({"status": "ok", "id": id}))
        }
        "done" => {
            let id = args
                .get("id")
                .and_then(|v| {
                    v.as_u64()
                        .or_else(|| v.as_str().and_then(|s| s.parse().ok()))
                })
                .map(|v| v as usize);
            match id {
                Some(id) => {
                    if ctx.mark_todo_done(id) {
                        Ok(json!({"status": "ok", "message": format!("Marked todo {id} as done")}))
                    } else {
                        Ok(json!({"error": format!("Todo {id} not found")}))
                    }
                }
                None => Ok(json!({"error": "id is required and must be a number"})),
            }
        }
        "list" => {
            let list = &ctx.todo_list;
            if list.is_empty() {
                Ok(json!({"goal": "", "todos": []}))
            } else {
                Ok(serde_json::to_value(list).unwrap_or(json!({"error": "serialization failed"})))
            }
        }
        "clear" => {
            ctx.clear_todo_list();
            Ok(json!({"status": "ok", "message": "Todo list cleared"}))
        }
        _ => bail!("Unknown todo action: {action}"),
    }
}