Skip to main content

archivist_core/
intent.rs

1//! Free-form mention handling — "talk to the Archivist and it does the thing."
2//!
3//! When the bot is @mentioned with a message that is NOT a `!`/slash command,
4//! this module classifies the intent (optionally via a local LLM) and maps it
5//! to a whitelisted core action. Everything the LLM produces is validated
6//! against a fixed whitelist before ANYTHING runs; prompt injection can only
7//! pick from the closed action set.
8//!
9//! LLM features are OPTIONAL: with `FANFIC_ARCHIVIST_LLM_ENABLED=0` (default)
10//! the module still handles fanfiction URLs (metadata/download/bookmark),
11//! questions ("how do I...") via `/api/docs/ask`, and everything else via a
12//! natural-language `/ask` fallback — no Ollama, no latency.
13
14use serde::{Deserialize, Serialize};
15
16use crate::cache::PageCache;
17use crate::config::BotConfig;
18use crate::util::{extract_fanfic_url, normalize_url};
19
20/// One parsed intent. The LLM can only produce actions in this enum.
21#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
22#[serde(tag = "action", rename_all = "snake_case")]
23pub enum Intent {
24    /// Search with optional filters.
25    Search {
26        /// Free-form query string.
27        query: Option<String>,
28        /// Fandom filter.
29        fandom: Option<String>,
30        /// Minimum word count.
31        min_words: Option<i64>,
32        /// Only completed works.
33        complete: Option<bool>,
34    },
35    /// Ask the archive a question (`/api/docs/ask`).
36    Ask {
37        /// The question.
38        query: String,
39    },
40    /// Find a passage by phrase (body search).
41    Quote {
42        /// Phrase to look for.
43        phrase: String,
44    },
45    /// Show recommendations for a fic.
46    Recs,
47    /// Fresh/recent works.
48    Fresh,
49    /// Hidden gems (underrated works).
50    Gems,
51    /// Random roll of a work.
52    Roll,
53    /// Download a fic in a format.
54    Download {
55        /// Story URL.
56        url: String,
57    },
58    /// Bookmark a fic.
59    Bookmark {
60        /// Story URL.
61        url: String,
62    },
63    /// Fetch metadata for a URL.
64    Metadata {
65        /// Story URL.
66        url: String,
67    },
68    /// Show help.
69    Help {
70        /// Optional help topic/question.
71        question: String,
72    },
73    /// No actionable intent (replies with a hint).
74    #[serde(rename = "none")]
75    None {
76        /// Fallback reply text.
77        reply: String,
78    },
79}
80
81impl Intent {
82    /// Human-readable one-liner for the "interpreted as" header.
83    pub fn describe(&self) -> String {
84        match self {
85            Self::Search { query, fandom, min_words, complete } => {
86                let mut s = String::from("search");
87                if let Some(q) = query {
88                    if !q.is_empty() {
89                        s.push_str(&format!(": \"{q}\""));
90                    }
91                }
92                if let Some(f) = fandom {
93                    if !f.is_empty() {
94                        s.push_str(&format!(" · fandom {f}"));
95                    }
96                }
97                if let Some(m) = min_words {
98                    s.push_str(&format!(" · ≥{m} words"));
99                }
100                if *complete == Some(true) {
101                    s.push_str(" · complete");
102                }
103                s
104            }
105            Self::Ask { query } => format!("ask: \"{query}\""),
106            Self::Quote { phrase } => format!("quote: \"{phrase}\""),
107            Self::Recs => "recommendations".into(),
108            Self::Fresh => "fresh recommendations".into(),
109            Self::Gems => "hidden gems".into(),
110            Self::Roll => "roll".into(),
111            Self::Download { url } => format!("download: {url}"),
112            Self::Bookmark { url } => format!("bookmark: {url}"),
113            Self::Metadata { url } => format!("metadata: {url}"),
114            Self::Help { question } => format!("help: {question}"),
115            Self::None { reply } => format!("chat: {}", crate::util::truncate(reply, 40)),
116        }
117    }
118}
119
120/// Normalize the message the way the LLM/cache sees it.
121pub fn cache_key_for(text: &str) -> String {
122    let t = text.trim().to_lowercase();
123    format!("archivist:intentcache:{}", sha256_hex(t.as_bytes()))
124}
125
126fn sha256_hex(bytes: &[u8]) -> String {
127    use sha2::{Digest, Sha256};
128    let mut hasher = Sha256::new();
129    hasher.update(bytes);
130    let out = hasher.finalize();
131    let mut s = String::with_capacity(64);
132    for b in out {
133        s.push_str(&format!("{b:02x}"));
134    }
135    s
136}
137
138/// The fixed whitelist of action names the LLM may emit. Anything else is
139/// rejected and treated as `None`.
140const ACTION_WHITELIST: &[&str] = &[
141    "search", "ask", "quote", "recs", "fresh", "gems", "roll", "metadata",
142    "download", "bookmark", "help", "none",
143];
144
145fn is_whitelisted(action: &str) -> bool {
146    ACTION_WHITELIST.contains(&action)
147}
148
149/// Validate + sanitize a parsed intent against the whitelist and URL rules.
150/// Returns the sanitized intent, or `None` if it must be dropped.
151fn sanitize(raw: &serde_json::Value) -> Option<Intent> {
152    let action = raw.get("action").and_then(|v| v.as_str()).unwrap_or("").to_string();
153    if !is_whitelisted(&action) {
154        return None;
155    }
156    let params = raw.get("params").cloned().unwrap_or_else(|| serde_json::json!({}));
157    let s = |k: &str| params.get(k).and_then(|v| v.as_str()).map(|v| v.to_string());
158    let i = |k: &str| params.get(k).and_then(|v| v.as_i64());
159    let b = |k: &str| params.get(k).and_then(|v| v.as_bool());
160
161    match action.as_str() {
162        "search" => {
163            let query = s("query").filter(|q| !q.is_empty());
164            let fandom = s("fandom").filter(|f| !f.is_empty());
165            let min_words = i("min_words").filter(|v| *v >= 0);
166            let complete = b("complete");
167            Some(Intent::Search { query, fandom, min_words, complete })
168        }
169        "ask" => s("query").filter(|q| !q.is_empty()).map(|query| Intent::Ask { query }),
170        "quote" => s("phrase").filter(|p| !p.is_empty()).map(|phrase| Intent::Quote { phrase }),
171        "recs" => Some(Intent::Recs),
172        "fresh" => Some(Intent::Fresh),
173        "gems" => Some(Intent::Gems),
174        "roll" => Some(Intent::Roll),
175        "download" => s("url").filter(|u| !u.is_empty()).map(|url| Intent::Download { url }),
176        "bookmark" => s("url").filter(|u| !u.is_empty()).map(|url| Intent::Bookmark { url }),
177        "metadata" => s("url").filter(|u| !u.is_empty()).map(|url| Intent::Metadata { url }),
178        "help" => s("question").filter(|q| !q.is_empty()).map(|question| Intent::Help { question }),
179        "none" => {
180            let reply = s("reply").unwrap_or_default();
181            Some(Intent::None { reply })
182        }
183        _ => None,
184    }
185}
186
187/// The prompt sent to the LLM (few-shot). Kept compact for token economy.
188fn prompt_for(text: &str) -> String {
189    format!(
190        r#"You convert a user's message into ONE action for the fanfic-archivist bot.
191Reply with ONLY a JSON object: {{"action": "...", "params": {{...}}}}. No prose.
192
193Actions (closed set — never invent others):
194- "search": find fics by keywords/filters. params: query(str), fandom(str?), min_words(int?), complete(bool?)
195- "ask": find fics from a natural-language description. params: query(str)
196- "quote": find fics containing an exact phrase/dialogue. params: phrase(str)
197- "recs" | "fresh" | "gems" | "roll": recommendation modes. params: {{}}
198- "download" | "bookmark" | "metadata": act on a fanfiction URL. params: url(str)
199- "help": a question about using the bot or FicHub. params: question(str)
200- "none": chit-chat / anything else. params: reply(str) — one short helpful line.
201
202Rules:
203- If the message contains a fanfiction URL, prefer metadata/download/bookmark.
204- If they describe what to read (trope/ship/fandom/length), use search or ask.
205- If they quote dialogue ("find the fic where X says Y"), use quote.
206- Keep params minimal and faithful. Do not invent filters.
207
208Examples:
209"completed dark harry over 50k" -> {{"action":"search","params":{{"query":"dark harry","complete":true,"min_words":50000}}}}
210"something cozy with a coffee shop au" -> {{"action":"ask","params":{{"query":"cozy coffee shop au"}}}}
211"where does she say 'I solemnly swear'" -> {{"action":"quote","params":{{"phrase":"I solemnly swear"}}}}
212"https://archiveofourown.org/works/123" -> {{"action":"metadata","params":{{"url":"https://archiveofourown.org/works/123"}}}}
213"how do I send to my kindle?" -> {{"action":"help","params":{{"question":"how do I send to my kindle?"}}}}
214"hello" -> {{"action":"none","params":{{"reply":"Hi! Ask me to find, download, or bookmark fics."}}}}
215
216User message: {text}
217"#
218    )
219}
220
221/// Classify a free-form message into an intent.
222///
223/// Order of operations:
224/// 1. URL in message (fanfiction host) → Metadata/Download/Bookmark immediately
225///    (no LLM needed, deterministic).
226/// 2. Cache hit → return cached intent.
227/// 3. LLM enabled → `POST /api/generate` (Ollama, format json) with timeout;
228///    validate against whitelist.
229/// 4. LLM disabled or failure → keyword heuristics, then natural-language `/ask`
230///    fallback handled by the caller.
231///
232/// Returns `(Intent, cached: bool)`.
233pub async fn classify(
234    text: &str,
235    config: &std::sync::Arc<BotConfig>,
236    cache: &PageCache,
237    http: &reqwest::Client,
238) -> (Intent, bool) {
239    let trimmed = text.trim();
240    if trimmed.is_empty() {
241        return (Intent::None { reply: "Hi! Ask me to find, download, or bookmark fics.".into() }, false);
242    }
243
244    // 1. URL fast-path.
245    if let Some(url) = extract_fanfic_url(trimmed) {
246        let normalized = normalize_url(&url);
247        let lower = trimmed.to_lowercase();
248        let intent = if lower.contains("bookmark") || lower.contains("book mark") {
249            Intent::Bookmark { url: normalized.clone() }
250        } else if lower.contains("download") {
251            Intent::Download { url: normalized }
252        } else {
253            Intent::Metadata { url: normalized }
254        };
255        return (intent, false);
256    }
257
258    // 2. Cache.
259    let key = cache_key_for(trimmed);
260    if let Some(v) = cache.cached_raw(&key).await {
261        if let Some(intent) = sanitize(&v) {
262            return (intent, true);
263        }
264    }
265
266    // 3. LLM path.
267    if config.llm_enabled {
268        if let Some(intent) = classify_with_llm(trimmed, config, http).await {
269            if let Ok(v) = serde_json::to_value(&intent) {
270                cache.cache_raw(&key, config.intent_cache_ttl, &v).await;
271            }
272            return (intent, false);
273        }
274    }
275
276    // 4. Heuristic fallback (no LLM or LLM failed).
277    (heuristic(trimmed), false)
278}
279
280/// Call Ollama `/api/generate` with `format:"json"` and parse the result.
281async fn classify_with_llm(
282    text: &str,
283    config: &std::sync::Arc<BotConfig>,
284    http: &reqwest::Client,
285) -> Option<Intent> {
286    let url = format!("{}/api/generate", config.ollama_url.trim_end_matches('/'));
287    let body = serde_json::json!({
288        "model": config.intent_model,
289        "prompt": prompt_for(text),
290        "stream": false,
291        "format": "json",
292        "options": { "temperature": 0 },
293    });
294    let resp = tokio::time::timeout(
295        std::time::Duration::from_secs(config.intent_timeout_secs),
296        http.post(&url).json(&body).send(),
297    )
298    .await
299    .ok()?
300    .ok()?;
301    let status = resp.status();
302    if !status.is_success() {
303        tracing::warn!("intent LLM HTTP {status}");
304        return None;
305    }
306    let json: serde_json::Value = resp.json().await.ok()?;
307    let text_out = json.get("response").and_then(|v| v.as_str())?;
308    let parsed: serde_json::Value = serde_json::from_str(text_out).ok()?;
309    sanitize(&parsed)
310}
311
312/// Keyword heuristic (deterministic; used when the LLM is off/failed).
313fn heuristic(text: &str) -> Intent {
314    let lower = text.to_lowercase();
315    let s = |words: &[&str]| words.iter().any(|w| lower.contains(w));
316
317    if s(&["recs", "recommend", "recommendation"]) {
318        return Intent::Recs;
319    }
320    if s(&["fresh"]) {
321        return Intent::Fresh;
322    }
323    if s(&["gems", "underrated"]) {
324        return Intent::Gems;
325    }
326    if s(&["roll", "random", "blind date", "surprise me"]) {
327        return Intent::Roll;
328    }
329    if s(&["how do i", "how do you", "how to", "what is", "what are", "help", "guide"]) {
330        return Intent::Help {
331            question: text.trim().to_string(),
332        };
333    }
334    if s(&["quote", "says", "said", "dialogue", "line where", "the fic where"]) {
335        return Intent::Quote {
336            phrase: text.trim().trim_start_matches("quote").trim().to_string(),
337        };
338    }
339    if s(&["bookmark", "save this", "add to library"]) {
340        return Intent::Ask {
341            query: text.trim().to_string(),
342        };
343    }
344    if s(&["download"]) {
345        return Intent::Ask {
346            query: text.trim().to_string(),
347        };
348    }
349    if s(&["search", "find", "looking for", "looking for a fic", "fic where", "fic about", "recommend me"]) {
350        return Intent::Ask {
351            query: text.trim().to_string(),
352        };
353    }
354    Intent::None {
355        reply: "Hi! Ask me to find, download, or bookmark fics.".to_string(),
356    }
357}
358
359#[cfg(test)]
360mod tests {
361    use super::*;
362
363    #[test]
364    fn sanitize_accepts_whitelisted() {
365        let v = serde_json::json!({
366            "action": "search",
367            "params": {"query": "dark harry", "complete": true, "min_words": 50000}
368        });
369        let i = sanitize(&v).unwrap();
370        assert_eq!(
371            i,
372            Intent::Search {
373                query: Some("dark harry".into()),
374                fandom: None,
375                min_words: Some(50000),
376                complete: Some(true),
377            }
378        );
379    }
380
381    #[test]
382    fn sanitize_rejects_unknown_action() {
383        let v = serde_json::json!({"action": "rate", "params": {"stars": 5}});
384        assert!(sanitize(&v).is_none());
385    }
386
387    #[test]
388    fn sanitize_rejects_injection_action() {
389        let v = serde_json::json!({"action": "admin", "params": {"cmd": "drop database"}});
390        assert!(sanitize(&v).is_none());
391    }
392
393    #[test]
394    fn sanitize_clamps_negative_words() {
395        let v = serde_json::json!({"action": "search", "params": {"query": "x", "min_words": -5}});
396        let i = sanitize(&v).unwrap();
397        match i {
398            Intent::Search { min_words, .. } => assert_eq!(min_words, None),
399            _ => panic!("expected search"),
400        }
401    }
402
403    #[test]
404    fn sanitize_requires_url_for_download() {
405        let v = serde_json::json!({"action": "download", "params": {}});
406        assert!(sanitize(&v).is_none());
407    }
408
409    #[test]
410    fn cache_key_is_stable_and_namespaced() {
411        let k1 = cache_key_for("  Find me a FIC  ");
412        let k2 = cache_key_for("find me a fic");
413        assert_eq!(k1, k2);
414        assert!(k1.starts_with("archivist:intentcache:"));
415        assert_eq!(k1.len(), "archivist:intentcache:".len() + 64);
416    }
417
418    #[test]
419    fn heuristic_detects_help() {
420        let i = heuristic("how do I send to my kindle?");
421        assert!(matches!(i, Intent::Help { .. }));
422    }
423
424    #[test]
425    fn heuristic_detects_recs() {
426        let i = heuristic("give me recommendations please");
427        assert_eq!(i, Intent::Recs);
428    }
429}