Newer
Older
TelosDB / src-backend / src / mcp / tools.rs
use crate::AppState;
use std::sync::Arc;
use super::handlers;

pub async fn handle_tools_list() -> anyhow::Result<serde_json::Value> {
    Ok(serde_json::json!({
        "tools": [
            {
                "name": "save_document",
                "description": "Save a document to the database",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "content": { "type": "string" },
                        "document_name": { "type": "string", "description": "Name or path of the document" }
                    },
                    "required": ["content", "document_name"]
                }
            },
            {
                "name": "find_documents",
                "description": "Find documents by text similarity",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "content": { "type": "string" },
                        "limit": { "type": "number", "description": "Maximum number of results to return (default 10)" }
                    },
                    "required": ["content"]
                }
            },
            {
                "name": "find_by_vector",
                "description": "Find documents by vector similarity",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "vector": { "type": "array", "items": { "type": "number" } },
                        "limit": { "type": "number", "description": "Maximum number of results to return (default 10)" }
                    },
                    "required": ["vector"]
                }
            },
            {
                "name": "delete_item",
                "description": "Delete an item by ID",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "id": { "type": "number" }
                    },
                    "required": ["id"]
                }
            },
            {
                "name": "get_vector",
                "description": "Get the vector embedding for an item by ID",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "id": { "type": "number" }
                    },
                    "required": ["id"]
                }
            },
            {
                "name": "get_document",
                "description": "Get the document content by ID",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "id": { "type": "number" }
                    },
                    "required": ["id"]
                }
            },
            {
                "name": "get_documents_count",
                "description": "Get the total number of documents",
                "inputSchema": {
                    "type": "object",
                    "properties": {}
                }
            },
            {
                "name": "list_documents",
                "description": "List documents with pagination (default order: oldest first)",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "limit": { "type": "number", "description": "Number of items to return (default 10)" },
                        "offset": { "type": "number", "description": "Number of items to skip (default 0)" }
                    }
                }
            },
            {
                "name": "llm_generate",
                "description": "Generate text via LLM",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "prompt": { "type": "string" },
                        "n_predict": { "type": "number" },
                        "temperature": { "type": "number" }
                    },
                    "required": ["prompt"]
                }
            },
            {
                "name": "read_recent_items",
                "description": "Read recent items from the database",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "limit": { "type": "number", "description": "Number of items to return (default 10)" }
                    }
                }
            }
        ]
    }))
}

pub async fn handle_tools_call(state: &Arc<AppState>, params: &serde_json::Value) -> anyhow::Result<serde_json::Value> {
    let tool_name = params["name"].as_str().unwrap_or("");
    let args = &params["arguments"];

    match tool_name {
        "save_document" => {
            let content = args["content"].as_str().ok_or(anyhow::anyhow!("content is required"))?;
            let document_name = args["document_name"].as_str().or(args["path"].as_str()).ok_or(anyhow::anyhow!("document_name is required"))?;
            handlers::handle_save_document(state, content, document_name).await
        }
        "find_documents" => {
            let content = args["content"].as_str().unwrap_or("");
            let limit = args["limit"].as_u64().unwrap_or(10) as usize;
            handlers::handle_find_documents(state, content, limit).await
        }
        "find_by_vector" => {
            let vector: Vec<f32> = args["vector"]
                .as_array()
                .unwrap_or(&vec![])
                .iter()
                .map(|v| v.as_f64().unwrap_or(0.0) as f32)
                .collect();
            let limit = args["limit"].as_u64().unwrap_or(10) as usize;
            handlers::handle_find_by_vector(state, vector, limit).await
        }
        "delete_item" => {
            let id = args["id"].as_i64().ok_or(anyhow::anyhow!("id is required"))? as i32;
            handlers::handle_delete_item(state, id).await
        }
        "get_vector" => {
            let id = args["id"].as_i64().ok_or(anyhow::anyhow!("id is required"))? as i32;
            handlers::handle_get_vector(state, id).await
        }
        "get_document" => {
            let id = args["id"].as_i64().ok_or(anyhow::anyhow!("id is required"))? as i32;
            handlers::handle_get_document(state, id).await
        }
        "get_documents_count" => {
            handlers::handle_get_documents_count(state).await
        }
        "list_documents" => {
            let limit = args["limit"].as_u64().unwrap_or(10) as u64;
            let offset = args["offset"].as_u64().unwrap_or(0) as u64;
            handlers::handle_list_documents(state, limit, offset).await
        }
        "llm_generate" => {
            let prompt = args["prompt"].as_str().unwrap_or("");
            let n_predict = args["n_predict"].as_i64().unwrap_or(128) as i32;
            let temperature = args["temperature"].as_f64().unwrap_or(0.7) as f32;
            handlers::handle_llm_generate(state, prompt, n_predict, temperature).await
        }
        "read_recent_items" => {
            let limit = args["limit"].as_u64().unwrap_or(10) as u64;
            handlers::handle_read_recent_items(state, limit).await
        }
        _ => Err(anyhow::anyhow!("Tool not found: {}", tool_name)),
    }
}