monothink 0.4.7

/*
* This file is part of monothink.
*
* monothink is free software: you can redistribute it and/or modify it under the terms
* of the GNU General Public License as published by the Free Software Foundation,
* either version 3 of the License, or (at your option) any later version.
*
* monothink is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
* See the GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License along with Wavething.
* If not, see <https://www.gnu.org/licenses/>.

*/

use std::{cmp, collections::BTreeMap};

use crate::{
    config::{self, McpTool, WebSearchProvider},
    llm::tools::ToolError,
};
use rig::tool::{Tool, ToolContext};
use rmcp::{
    RoleClient,
    model::{CallToolRequestParams, CallToolResult, ClientInfo},
    service::RunningService,
};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};

const MAX_RESULT_LIMIT: u64 = 20;

#[derive(Deserialize, Serialize)]
pub struct WebSearchArgs {
    query: String,
    #[serde(default = "web_search_max_results_default")]
    max_results: u64,
    // possible other parameters:
    // search_depth
    // start_date
    // end_date
    // exclude_domains
    // include_domains
    //
}

fn web_search_max_results_default() -> u64 {
    5
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct WebSearchResult {
    url: String,
    title: String,
    content: String,
}

#[derive(Debug)]
enum WebSearchExecutor {
    SearXNG {
        client: reqwest::Client,
        url: String,
    },
    Tavily {
        provider: WebSearchProvider,
        service: RunningService<RoleClient, ClientInfo>,
    },
}

impl WebSearchExecutor {
    // @todo take channel to report warning messages in case of errors
    async fn new(config: &config::WebSearch) -> Option<WebSearchExecutor> {
        match config.provider {
            WebSearchProvider::SearXNG => Self::build_rest_web_search(config).await,
            WebSearchProvider::Tavily => Self::build_mcp_web_search(config).await,
        }
    }

    async fn build_mcp_web_search(config: &config::WebSearch) -> Option<WebSearchExecutor> {
        let search_name = match config.provider.search_name() {
            Some(name) => name,
            None => return None,
        };

        let tool = super::get_mcp_tools(&BTreeMap::from([(
            config.provider.display_string(),
            McpTool {
                url: config.url.clone(),
            },
        )]))
        .await;
        // @todo check success

        // @todo is there an easier way to consume the vector?
        if let Some((service, tools)) = tool.into_iter().next() {
            // @todo validate schema (the parts we are interested in)
            // for default parameters, issue a warning only
            // for mandatory paremters, issue a warning and fail
            //tool.with_input_schema()
            // tool.output_schema (if present)
            tools
                .iter()
                // cheap way of ensuring the tool is there
                .find(|tool| tool.name == search_name)
                .map(|_tool| Self::Tavily {
                    provider: config.provider,
                    service,
                })
        } else {
            None
        }
    }

    async fn build_rest_web_search(config: &config::WebSearch) -> Option<WebSearchExecutor> {
        Some(WebSearchExecutor::SearXNG {
            client: reqwest::Client::new(),
            url: config.url.clone(),
        })
    }

    async fn call(&self, args: WebSearchArgs) -> Result<Vec<WebSearchResult>, ToolError> {
        match self {
            WebSearchExecutor::SearXNG { client, url } => {
                let params = Self::prepare_searxng_parameters(args);
                let response = Self::call_rest(client, url, &params).await?;
                Self::map_searxng_result(response)
            }
            WebSearchExecutor::Tavily { provider, service } => {
                let response = Self::call_mcp(provider, service, args).await?;
                Self::map_tavily_result(response)
            }
        }
    }

    async fn call_mcp(
        provider: &WebSearchProvider,
        service: &RunningService<RoleClient, ClientInfo>,
        args: WebSearchArgs,
        // ) -> Result<Vec<WebSearchResult>, ToolError> {
    ) -> Result<CallToolResult, ToolError> {
        let search_name = provider
            .search_name()
            .expect("Web search name is required for construction and must be present.");

        let arguments = serde_json::to_value(&args)
            .unwrap()
            .as_object()
            .cloned()
            .unwrap(); // @todo proper error handling

        let params = CallToolRequestParams::new(search_name).with_arguments(arguments);
        service.call_tool(params).await.map_err(|err| match err {
            rmcp::ServiceError::McpError(error_data) => ToolError::Mcp {
                kind: "MCP".to_string(),
                info: error_data.to_string(),
            },
            rmcp::ServiceError::TransportSend(dynamic_transport_error) => ToolError::Mcp {
                kind: "Transport send".to_string(),
                info: dynamic_transport_error.to_string(),
            },
            rmcp::ServiceError::TransportClosed => ToolError::Mcp {
                kind: "Transport closed".to_string(),
                info: String::new(),
            },
            rmcp::ServiceError::UnexpectedResponse => ToolError::Mcp {
                kind: "Unexpected response".to_string(),
                info: String::new(),
            },
            rmcp::ServiceError::Cancelled { reason } => ToolError::Mcp {
                kind: "Cancelled".to_string(),
                info: reason.unwrap_or_default(),
            },
            rmcp::ServiceError::Timeout { timeout } => ToolError::Mcp {
                kind: "Cancelled".to_string(),
                info: String::new(),
            },
            _ => ToolError::Mcp {
                kind: "Unknown".to_string(),
                info: String::new(),
            },
        })
    }

    async fn call_rest(
        client: &reqwest::Client,
        url: &str,
        params: &[(String, String)],
    ) -> Result<Value, ToolError> {
        // @todo this should be done upfront in prepare_searxng_args()
        client
            .get(url)
            .query(&params)
            .send()
            .await
            .map_err(|e| ToolError::NetworkError {
                reason: format!("{}", e),
            })?
            .json()
            .await
            .map_err(|e| ToolError::NetworkError {
                reason: format!("{}", e),
            })
    }

    fn map_searxng_result(response: Value) -> Result<Vec<WebSearchResult>, ToolError> {
        Ok(response
            .get("results")
            .expect("Response should contain results")
            .as_array()
            .expect("Schema does not match for results array")
            .iter()
            .map(|item| {
                let url = item
                    .get("url")
                    .expect("Schema does not match for url")
                    .as_str()
                    .unwrap()
                    .to_owned();
                let title = item
                    .get("title")
                    .expect("Schema does not match for title")
                    .as_str()
                    .unwrap()
                    .to_owned();
                let content = item
                    .get("content")
                    .expect("Schema does not match for content")
                    .as_str()
                    .unwrap()
                    .to_owned();
                WebSearchResult {
                    url,
                    title,
                    content,
                }
            })
            .collect())
    }

    fn map_tavily_result(
        result: rmcp::model::CallToolResult,
    ) -> Result<Vec<WebSearchResult>, ToolError> {
        // @todo better error handling: drop elements that fail and log a warning
        // we should also let the user and the LLM know that there is a technical error and disable the tool
        Ok(result
            .content
            .iter()
            .filter_map(|item| item.as_text())
            .flat_map(|item| {
                let result: Value =
                    serde_json::from_str(&item.text).expect("Response should be parsable as JSON");
                let results = result
                    .get("results")
                    .expect("Schema does not match for results")
                    .as_array()
                    .expect("Schema does not match for results array");
                results
                    .iter()
                    .map(|item| {
                        let url = item
                            .get("url")
                            .expect("Schema does not match for url")
                            .as_str()
                            .unwrap()
                            .to_owned();
                        let title = item
                            .get("title")
                            .expect("Schema does not match for title")
                            .as_str()
                            .unwrap()
                            .to_owned();
                        let content = item
                            .get("content")
                            .expect("Schema does not match for content")
                            .as_str()
                            .unwrap()
                            .to_owned();
                        WebSearchResult {
                            url,
                            title,
                            content,
                        }
                    })
                    .collect::<Vec<_>>()
            })
            .collect())
    }

    fn prepare_searxng_parameters(args: WebSearchArgs) -> Vec<(String, String)> {
        let limit = format!("{}", args.max_results);
        vec![
            ("q".to_string(), args.query),
            ("format".to_string(), "json".to_string()),
            ("limit".to_string(), limit),
        ]
    }
}

// @todo allow default parameters: check schema during new (issue a warning if not found) and apply them during call
#[derive(Debug)]
pub struct WebSearch {
    exec: WebSearchExecutor,
}

impl WebSearch {
    pub async fn new(config: &config::WebSearch) -> Option<Self> {
        WebSearchExecutor::new(config)
            .await
            .map(|exec| WebSearch { exec })
    }
}

impl Tool for WebSearch {
    const NAME: &'static str = "web_search";

    type Error = ToolError;
    type Args = WebSearchArgs;
    type Output = Vec<WebSearchResult>;

    async fn call(
        &self,
        _context: &mut ToolContext,
        args: Self::Args,
    ) -> Result<Self::Output, Self::Error> {
        let args = Self::Args {
            max_results: cmp::min(args.max_results, MAX_RESULT_LIMIT),
            ..args
        };
        self.exec.call(args).await
    }

    fn description(&self) -> String {
        "Search the internet for up-to-date information, recent events, live data, \
        or obscure facts not covered in your training data. \
        Use when the user asks for current news, real-time statuses, \
        or when you need to verify uncertain information. \
        Do not use for common knowledge or questions answerable with local tools."
            .to_string()
    }

    // could be extracted from schema but I don't like their description
    fn parameters(&self) -> serde_json::Value {
        let max_results_description = format!(
            "Maximum number of results to return. Recommended: 3-5. \
                    Higher values increase latency and token usage but may improve coverage. \
                    Defaults to 5 if omitted. Maximum value is {}.",
            MAX_RESULT_LIMIT
        );
        json!({
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "A concise, specific search query. Use keywords or short phrases. \
                    Avoid full sentences or overly broad terms for better results."
                },
                "max_results": {
                    "type": "integer",
                    "description": max_results_description,
                    "default": 5
                },
            },
            "required": ["query"]
        })
    }

    fn map_error(&self, error: Self::Error) -> rig::tool::ToolExecutionError {
        rig::tool::ToolExecutionError::from_error(error)
    }
}

// implementation non-config parts
impl WebSearchProvider {
    pub fn search_name(&self) -> Option<String> {
        match self {
            WebSearchProvider::SearXNG => Some("/search".to_string()), // REST endpoint
            WebSearchProvider::Tavily => Some("tavily_search".to_string()),
        }
    }
}

// @todo real tests check that new() works (schema is valid etc)
// @todo real tests check that call() returns successfull (ignoring content)
#[tokio::test]
async fn test_web_search() {
    // fake test to just see the serch results
    let config = config::load_config().expect("Should be able to load config.");
    let ws = WebSearch::new(&config::WebSearch {
        provider: WebSearchProvider::Tavily,
        url: config
            .tools
            .web_search
            .expect("Should have websearch configured for tests")
            .url,
    })
    .await
    .expect("Should be able to instantiate web search");
    let result = ws
        .call(
            &mut ToolContext::default(),
            WebSearchArgs {
                query: "Rust latest version".to_string(),
                max_results: 5,
            },
        )
        .await
        .unwrap();
    println!("Result: {:?}", result);
}