ratatoskr/
request.rs

1//! What a search asks for, beyond the words: how many results, where, how
2//! recent, how safe, which page. Every engine reads what it supports and
3//! ignores the rest (an encyclopedia has no "past week").
4
5use std::fmt;
6use std::str::FromStr;
7
8/// How strictly adult content is filtered.
9#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
10pub enum SafeSearch {
11    On,
12    #[default]
13    Moderate,
14    Off,
15}
16
17/// How recent results must be.
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
19pub enum TimeLimit {
20    Day,
21    Week,
22    Month,
23    Year,
24}
25
26impl TimeLimit {
27    /// The one-letter code most engines use: d, w, m, y.
28    pub fn letter(self) -> &'static str {
29        match self {
30            Self::Day => "d",
31            Self::Week => "w",
32            Self::Month => "m",
33            Self::Year => "y",
34        }
35    }
36}
37
38/// An option value that is not one of the allowed ones.
39#[derive(Clone, Debug, PartialEq, Eq)]
40pub struct BadOption {
41    pub option: &'static str,
42    pub value: String,
43    pub allowed: &'static str,
44}
45
46impl fmt::Display for BadOption {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        write!(f, "{} must be one of {}, not {:?}", self.option, self.allowed, self.value)
49    }
50}
51
52impl std::error::Error for BadOption {}
53
54impl FromStr for SafeSearch {
55    type Err = BadOption;
56    fn from_str(value: &str) -> Result<Self, BadOption> {
57        match value.trim().to_ascii_lowercase().as_str() {
58            "on" | "strict" => Ok(Self::On),
59            "moderate" => Ok(Self::Moderate),
60            "off" => Ok(Self::Off),
61            _ => Err(BadOption { option: "safesearch", value: value.into(), allowed: "on, moderate, off" }),
62        }
63    }
64}
65
66impl FromStr for TimeLimit {
67    type Err = BadOption;
68    fn from_str(value: &str) -> Result<Self, BadOption> {
69        match value.trim().to_ascii_lowercase().as_str() {
70            "d" | "day" => Ok(Self::Day),
71            "w" | "week" => Ok(Self::Week),
72            "m" | "month" => Ok(Self::Month),
73            "y" | "year" => Ok(Self::Year),
74            _ => Err(BadOption { option: "time", value: value.into(), allowed: "day, week, month, year" }),
75        }
76    }
77}
78
79/// Why JSON arguments are not a search.
80#[derive(Clone, Debug, PartialEq, Eq)]
81pub enum RequestError {
82    NoQuery,
83    Bad(BadOption),
84}
85
86impl fmt::Display for RequestError {
87    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88        match self {
89            Self::NoQuery => write!(f, "a query is needed"),
90            Self::Bad(bad) => write!(f, "{bad}"),
91        }
92    }
93}
94
95impl std::error::Error for RequestError {}
96
97/// A search: the words and how to search for them.
98#[derive(Clone, Debug, PartialEq, Eq)]
99pub struct SearchRequest {
100    pub query: String,
101    /// At most this many results.
102    pub limit: usize,
103    /// Localize to a country: an ISO 3166-1 alpha-2 code ("US", "IN", "DE").
104    pub country: Option<String>,
105    pub safesearch: SafeSearch,
106    pub time: Option<TimeLimit>,
107    /// The page of results, from 1.
108    pub page: u32,
109}
110
111impl SearchRequest {
112    pub fn new(query: impl Into<String>, limit: usize) -> Self {
113        Self { query: query.into(), limit, country: None, safesearch: SafeSearch::default(), time: None, page: 1 }
114    }
115
116    pub fn country(mut self, country: Option<&str>) -> Self {
117        self.country = country.map(str::trim).filter(|c| !c.is_empty()).map(str::to_ascii_uppercase);
118        self
119    }
120
121    pub fn safesearch(mut self, safesearch: SafeSearch) -> Self {
122        self.safesearch = safesearch;
123        self
124    }
125
126    pub fn time(mut self, time: Option<TimeLimit>) -> Self {
127        self.time = time;
128        self
129    }
130
131    pub fn page(mut self, page: u32) -> Self {
132        self.page = page.max(1);
133        self
134    }
135
136    /// The same request for `limit` results.
137    pub fn with_limit(&self, limit: usize) -> Self {
138        Self { limit, ..self.clone() }
139    }
140
141    /// Whether only the first page, at any date, is asked for: the plain
142    /// search an encyclopedia's one best page belongs to.
143    pub fn is_plain(&self) -> bool {
144        self.page == 1 && self.time.is_none()
145    }
146
147    /// A request from JSON arguments, as the MCP tool and the HTTP API take
148    /// them: `query` (required), `limit` (default `default_limit`, at most
149    /// `max_limit`), `country`, `time`, `safesearch` and `page`.
150    pub fn from_json(arguments: &serde_json::Value, default_limit: u64, max_limit: u64) -> Result<Self, RequestError> {
151        use serde_json::Value;
152        let text = |key: &str| arguments.get(key).and_then(Value::as_str).map(str::trim).filter(|v| !v.is_empty());
153        let query = text("query").ok_or(RequestError::NoQuery)?;
154        let limit = arguments.get("limit").and_then(Value::as_u64).unwrap_or(default_limit).clamp(1, max_limit);
155        let time = text("time").map(str::parse::<TimeLimit>).transpose().map_err(RequestError::Bad)?;
156        let safesearch = text("safesearch").map(str::parse::<SafeSearch>).transpose().map_err(RequestError::Bad)?;
157        let page = arguments.get("page").and_then(Value::as_u64).unwrap_or(1).clamp(1, 100) as u32;
158        Ok(Self::new(query, limit as usize)
159            .country(text("country"))
160            .time(time)
161            .safesearch(safesearch.unwrap_or_default())
162            .page(page))
163    }
164
165    /// The key a cached answer to this request is stored under.
166    pub fn cache_key(&self, category: &str) -> String {
167        format!(
168            "{category}\u{0}{}\u{0}{}\u{0}{:?}\u{0}{:?}\u{0}{}\u{0}{}",
169            self.limit,
170            self.country.as_deref().unwrap_or(""),
171            self.safesearch,
172            self.time,
173            self.page,
174            self.query.trim().to_lowercase()
175        )
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182
183    #[test]
184    fn options_parse_and_normalize() {
185        assert_eq!("Week".parse::<TimeLimit>().unwrap(), TimeLimit::Week);
186        assert_eq!("d".parse::<TimeLimit>().unwrap().letter(), "d");
187        assert_eq!("strict".parse::<SafeSearch>().unwrap(), SafeSearch::On);
188        let error = "fortnight".parse::<TimeLimit>().unwrap_err();
189        assert_eq!(error.to_string(), "time must be one of day, week, month, year, not \"fortnight\"");
190        let request = SearchRequest::new("Rust", 5).country(Some(" in ")).page(0).time(Some(TimeLimit::Day));
191        assert_eq!((request.country.as_deref(), request.page), (Some("IN"), 1));
192        assert!(!request.is_plain() && SearchRequest::new("x", 5).is_plain());
193        assert_ne!(
194            request.cache_key("text"),
195            request.clone().time(None).cache_key("text"),
196            "options make their own cache entry"
197        );
198        assert_ne!(request.cache_key("text"), request.cache_key("news"));
199    }
200}
ratatoskr/request.rs

ratatoskr/
request.rs
