bot/src/domain/agent/tools.rs

use super::api::{self, FunctionDefinition, Message, Tool, ToolCall};
use sea_orm::{
    ColumnTrait, Condition, DatabaseConnection, EntityTrait, QueryFilter, QueryOrder, QuerySelect,
};
use serde::Deserialize;
use uuid::Uuid;

#[derive(Deserialize)]
struct GoogleSearchArgs {
    query: String,
}

#[derive(Deserialize)]
struct FinishArgs {
    result: String,
}

#[derive(Deserialize)]
struct ListRunsArgs {
    from: Option<String>,
    to: Option<String>,
    task_ids: Option<Vec<String>>,
}

pub fn get_tools() -> Vec<Tool> {
    vec![
        Tool {
            tool_type: "function".to_string(),
            function: FunctionDefinition {
                name: "google_search".to_string(),
                description: "Search the web for information".to_string(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query"
                        }
                    },
                    "required": ["query"]
                }),
            },
        },
        Tool {
            tool_type: "function".to_string(),
            function: FunctionDefinition {
                name: "finish".to_string(),
                description: "Finish the task and provide a final answer".to_string(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "result": {
                            "type": "string",
                            "description": "The final detailed answer to the task"
                        }
                    },
                    "required": ["result"]
                }),
            },
        },
        Tool {
            tool_type: "function".to_string(),
            function: FunctionDefinition {
                name: "list_tasks".to_string(),
                description: "List all existing tasks and their goals".to_string(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {},
                }),
            },
        },
        Tool {
            tool_type: "function".to_string(),
            function: FunctionDefinition {
                name: "list_runs".to_string(),
                description: "List task runs with optional filters for date and task IDs"
                    .to_string(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "from": {
                            "type": "string",
                            "description": "ISO 8601 date string to filter runs from"
                        },
                        "to": {
                            "type": "string",
                            "description": "ISO 8601 date string to filter runs to"
                        },
                        "task_ids": {
                            "type": "array",
                            "items": { "type": "string", "format": "uuid" },
                            "description": "List of task IDs to filter runs for"
                        },

                    },
                }),
            },
        },
    ]
}

pub async fn handle_tool_call(
    tool_call: &ToolCall,
    tavily_api_key: &Option<String>,
    db: &DatabaseConnection,
) -> Result<(Message, bool, Option<String>), Box<dyn std::error::Error>> {
    let mut answer = None;
    let name = &tool_call.function.name;

    let (content, written) = if name == "google_search" {
        let args: GoogleSearchArgs = serde_json::from_str(&tool_call.function.arguments)?;
        let query = &args.query;

        let search_result = if let Some(key) = tavily_api_key {
            match api::perform_search(query, key).await {
                Ok(results) => results,
                Err(e) => format!("Search error: {}", e),
            }
        } else {
            "Error: TAVILY_API_KEY is not set. Cannot perform real search.".to_string()
        };
        (search_result, false)
    } else if name == "finish" {
        let args: FinishArgs = serde_json::from_str(&tool_call.function.arguments)?;
        let result = &args.result;

        answer = Some(result.clone());
        (result.clone(), true)
    } else if name == "list_tasks" {
        tracing::info!("Listing tasks from database");
        use crate::entities::task;
        let tasks = task::Entity::find()
            .order_by_desc(task::Column::CreatedAt)
            .all(db)
            .await?;

        let mut out = String::from("Tasks:\n");
        for t in tasks {
            out.push_str(&format!("- ID: {}, Goal: {}\n", t.id, t.goal));
        }
        (out, false)
    } else if name == "list_runs" {
        use crate::entities::task_run;
        tracing::info!(
            "Listing runs from database, args: {}",
            tool_call.function.arguments
        );
        let args: ListRunsArgs = serde_json::from_str(&tool_call.function.arguments)?;

        tracing::info!(
            "Listing runs from database with filters: from={:?}, to={:?}, task_ids={:?}",
            args.from,
            args.to,
            args.task_ids
        );

        let query = task_run::Entity::find();

        let mut condition = Condition::all();

        if let Some(from_str) = args.from {
            if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(&from_str) {
                condition = condition.add(task_run::Column::CreatedAt.gte(dt));
            }
        }

        if let Some(to_str) = args.to {
            if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(&to_str) {
                condition = condition.add(task_run::Column::CreatedAt.lte(dt));
            }
        }

        if let Some(id_strs) = args.task_ids {
            let mut valid_ids = Vec::new();
            for id_str in id_strs {
                match id_str.parse::<Uuid>() {
                    Ok(uuid) => valid_ids.push(uuid),
                    Err(e) => tracing::warn!("Skipping invalid UUID '{}' from LLM: {}", id_str, e),
                }
            }
            if !valid_ids.is_empty() {
                condition = condition.add(task_run::Column::TaskId.is_in(valid_ids));
            }
        }

        let runs = query
            .filter(condition)
            .order_by_desc(task_run::Column::CreatedAt)
            .limit(20)
            .all(db)
            .await?;

        let mut out = String::from("Recent Runs:\n");
        for r in runs {
            out.push_str(&format!(
                "- ID: {}, Task ID: {}, Status: {}, Created At: {}, Answer: {:?}\n",
                r.id, r.task_id, r.status, r.created_at, r.answer
            ));
        }
        (out, false)
    } else {
        (format!("Error: Unknown tool {}", name), false)
    };

    Ok((
        Message {
            role: "tool".to_string(),
            content: Some(content),
            tool_calls: None,
            tool_call_id: Some(tool_call.id.clone()),
        },
        written,
        answer,
    ))
}