sol/src/tools/mod.rs

pub mod room_history;
pub mod room_info;
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;

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"]
                }),
            ),
        ]
    }

    pub async fn execute(&self, name: &str, arguments: &str) -> 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,
            _ => anyhow::bail!("Unknown tool: {name}"),
        }
    }
}