sol/src/tools/mod.rs

pub mod room_history;
pub mod room_info;
pub mod script;
pub mod search;

use std::sync::Arc;

use matrix_sdk::Client as MatrixClient;
use mistralai_client::v1::tool::Tool;
use opensearch::OpenSearch;
use serde_json::json;

use crate::config::Config;
use crate::context::ResponseContext;

pub struct ToolRegistry {
    opensearch: OpenSearch,
    matrix: MatrixClient,
    config: Arc<Config>,
}

impl ToolRegistry {
    pub fn new(opensearch: OpenSearch, matrix: MatrixClient, config: Arc<Config>) -> Self {
        Self {
            opensearch,
            matrix,
            config,
        }
    }

    pub fn tool_definitions() -> Vec<Tool> {
        vec![
            Tool::new(
                "search_archive".into(),
                "Search the message archive. Use this to find past conversations, \
                 messages from specific people, or about specific topics."
                    .into(),
                json!({
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Search query for message content"
                        },
                        "room": {
                            "type": "string",
                            "description": "Filter by room name (optional)"
                        },
                        "sender": {
                            "type": "string",
                            "description": "Filter by sender display name (optional)"
                        },
                        "after": {
                            "type": "string",
                            "description": "Unix timestamp in ms — only messages after this time (optional)"
                        },
                        "before": {
                            "type": "string",
                            "description": "Unix timestamp in ms — only messages before this time (optional)"
                        },
                        "limit": {
                            "type": "integer",
                            "description": "Max results to return (default 10)"
                        },
                        "semantic": {
                            "type": "boolean",
                            "description": "Use semantic search instead of keyword (optional)"
                        }
                    },
                    "required": ["query"]
                }),
            ),
            Tool::new(
                "get_room_context".into(),
                "Get messages around a specific point in time or event in a room. \
                 Useful for understanding the context of a conversation."
                    .into(),
                json!({
                    "type": "object",
                    "properties": {
                        "room_id": {
                            "type": "string",
                            "description": "The Matrix room ID"
                        },
                        "around_timestamp": {
                            "type": "integer",
                            "description": "Unix timestamp in ms to center the context around"
                        },
                        "around_event_id": {
                            "type": "string",
                            "description": "Event ID to center the context around"
                        },
                        "before_count": {
                            "type": "integer",
                            "description": "Number of messages before the pivot (default 10)"
                        },
                        "after_count": {
                            "type": "integer",
                            "description": "Number of messages after the pivot (default 10)"
                        }
                    },
                    "required": ["room_id"]
                }),
            ),
            Tool::new(
                "list_rooms".into(),
                "List all rooms Sol is currently in, with names and member counts.".into(),
                json!({
                    "type": "object",
                    "properties": {}
                }),
            ),
            Tool::new(
                "get_room_members".into(),
                "Get the list of members in a specific room.".into(),
                json!({
                    "type": "object",
                    "properties": {
                        "room_id": {
                            "type": "string",
                            "description": "The Matrix room ID"
                        }
                    },
                    "required": ["room_id"]
                }),
            ),
            Tool::new(
                "run_script".into(),
                "Execute a TypeScript/JavaScript snippet in a sandboxed runtime. \
                 Use this for math, date calculations, data transformations, or any \
                 computation that needs precision. The script has access to:\n\
                 - sol.search(query, opts?) — search the message archive. opts: \
                   { room?, sender?, after?, before?, limit?, semantic? }\n\
                 - sol.rooms() — list joined rooms (returns array of {name, id, members})\n\
                 - sol.members(roomName) — get room members (returns array of {name, id})\n\
                 - sol.fetch(url) — HTTP GET (allowlisted domains only)\n\
                 - sol.memory.get(query?) — retrieve internal notes relevant to the query\n\
                 - sol.memory.set(content, category?) — save an internal note for later reference\n\
                 - sol.fs.read(path), sol.fs.write(path, content), sol.fs.list(path?) — \
                   sandboxed temp filesystem for intermediate files\n\
                 - console.log() to produce output\n\
                 All sol.* methods are async — use await. The last expression value is \
                 also captured. Output is truncated to 4096 chars."
                    .into(),
                json!({
                    "type": "object",
                    "properties": {
                        "code": {
                            "type": "string",
                            "description": "TypeScript or JavaScript code to execute"
                        }
                    },
                    "required": ["code"]
                }),
            ),
        ]
    }

    pub async fn execute(
        &self,
        name: &str,
        arguments: &str,
        response_ctx: &ResponseContext,
    ) -> anyhow::Result<String> {
        match name {
            "search_archive" => {
                search::search_archive(
                    &self.opensearch,
                    &self.config.opensearch.index,
                    arguments,
                )
                .await
            }
            "get_room_context" => {
                room_history::get_room_context(
                    &self.opensearch,
                    &self.config.opensearch.index,
                    arguments,
                )
                .await
            }
            "list_rooms" => room_info::list_rooms(&self.matrix).await,
            "get_room_members" => room_info::get_room_members(&self.matrix, arguments).await,
            "run_script" => {
                script::run_script(
                    &self.opensearch,
                    &self.matrix,
                    &self.config,
                    arguments,
                    response_ctx,
                )
                .await
            }
            _ => anyhow::bail!("Unknown tool: {name}"),
        }
    }
}