openbrain-mcp/src/tools/mod.rs

//! MCP Tools for OpenBrain

pub mod batch_store;
pub mod query;
pub mod store;
pub mod purge;

use anyhow::Result;
use serde_json::{json, Value};
use std::sync::Arc;

use crate::AppState;

pub fn get_tool_definitions() -> Vec<Value> {
    vec![
        json!({
            "name": "store",
            "description": "Store a memory with automatic embedding generation and keyword extraction. The memory will be associated with the agent_id for isolated retrieval.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "content": {
                        "type": "string",
                        "description": "The text content to store as a memory"
                    },
                    "agent_id": {
                        "type": "string",
                        "description": "Unique identifier for the agent storing the memory (default: 'default')"
                    },
                    "metadata": {
                        "type": "object",
                        "description": "Optional metadata to attach to the memory"
                    },
                    "ttl": {
                        "type": "string",
                        "description": "Optional time-to-live for transient facts, like 30s, 15m, 1h, 7d, or 2w"
                    }
                },
                "required": ["content"]
            }
        }),
        json!({
            "name": "batch_store",
            "description": "Store multiple memories with automatic embedding generation and keyword extraction. Accepts 1-50 entries and stores them atomically in a single transaction.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "agent_id": {
                        "type": "string",
                        "description": "Unique identifier for the agent storing the memories (default: 'default')"
                    },
                    "ttl": {
                        "type": "string",
                        "description": "Optional default time-to-live applied to entries without their own ttl"
                    },
                    "entries": {
                        "type": "array",
                        "description": "Array of 1-50 memory entries to store atomically",
                        "items": {
                            "type": "object",
                            "properties": {
                                "content": {
                                    "type": "string",
                                    "description": "The text content to store as a memory"
                                },
                                "metadata": {
                                    "type": "object",
                                    "description": "Optional metadata to attach to the memory"
                                },
                                "ttl": {
                                    "type": "string",
                                    "description": "Optional per-entry time-to-live override like 30s, 15m, 1h, 7d, or 2w"
                                }
                            },
                            "required": ["content"]
                        }
                    }
                },
                "required": ["entries"]
            }
        }),
        json!({
            "name": "query",
            "description": "Query stored memories using hybrid semantic plus keyword search. Returns the most relevant memories based on vector similarity and exact-text ranking.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query text"
                    },
                    "agent_id": {
                        "type": "string",
                        "description": "Agent ID to search within (default: 'default')"
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum number of results to return (default: 10)"
                    },
                    "threshold": {
                        "type": "number",
                        "description": "Minimum similarity threshold 0.0-1.0 (default: 0.5)"
                    }
                },
                "required": ["query"]
            }
        }),
        json!({
            "name": "purge",
            "description": "Delete memories for an agent. Can delete all memories or those before a specific timestamp.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "agent_id": {
                        "type": "string",
                        "description": "Agent ID whose memories to delete (required)"
                    },
                    "before": {
                        "type": "string",
                        "description": "Optional ISO8601 timestamp - delete memories created before this time"
                    },
                    "confirm": {
                        "type": "boolean",
                        "description": "Must be true to confirm deletion"
                    }
                },
                "required": ["agent_id", "confirm"]
            }
        })
    ]
}

pub async fn execute_tool(
    state: &Arc<AppState>,
    tool_name: &str,
    arguments: Value,
) -> Result<String> {
    match tool_name {
        "store" => store::execute(state, arguments).await,
        "batch_store" => batch_store::execute(state, arguments).await,
        "query" => query::execute(state, arguments).await,
        "purge" => purge::execute(state, arguments).await,
        _ => anyhow::bail!("Unknown tool: {}", tool_name),
    }
}