1//! Choosing the picture for a fact: shortlist by meaning, then let Jev choose.
2//!
3//! The pipeline is pure over two ports and a cache. The embedder turns the names of the pictures (once, kept in
4//! [`IconVectors`]) and the fact's words into vectors; the closest [`SHORTLIST`] names go to Jev as one Choice,
5//! which may answer "none fits". Nothing here is hosted anywhere: an adapter gives the embedder and the judge and
6//! holds the cache. A fact is given a picture of a *thing*, never one of the interface icons
7//! (`IconId::depicts_a_thing`).
8
9use std::fmt;
10use std::sync::OnceLock;
11
12use ::log::{debug, info, warn};
13use whiskers_core::{IconId, cosine};
14
15use crate::embed::{EMBED_CHARS_MAX, EMBED_TEXTS_MAX, EmbedBatch, EmbedError, Embedder};
16use crate::judge::{Judge, JudgeError};
17
18/// How many candidates go to Jev. Jev answers thirty questions a minute and a Choice carries every option, so
19/// one question covers the whole shortlist; twenty keeps the question short enough to be answered well.
20pub const SHORTLIST: usize = 20;
21
22/// What an icon is called when it is turned into a vector, and what Jev is told it shows: the name read as words
23/// (`birthday-cake` is a picture of birthday cake).
24pub fn phrase(icon: &IconId) -> String {
25    format!("a picture of {}", icon.as_str().replace('-', " "))
26}
27
28/// The admitted pictures of things and their vectors, computed once and kept. A cache, not state anyone relies
29/// on: an adapter that holds one per process loses nothing when it restarts but a few embeddings' time.
30#[derive(Default)]
31pub struct IconVectors(OnceLock<Vec<(IconId, Vec<f32>)>>);
32
33impl IconVectors {
34    /// The vectors, computing them with `embedder` the first time.
35    pub async fn get(&self, embedder: &impl Embedder) -> Result<&[(IconId, Vec<f32>)], EmbedError> {
36        if let Some(v) = self.0.get() {
37            return Ok(v);
38        }
39        let icons: Vec<IconId> = IconId::all().filter(IconId::depicts_a_thing).collect();
40        let mut out = Vec::with_capacity(icons.len());
41        for chunk in icons.chunks(EMBED_TEXTS_MAX) {
42            let batch = EmbedBatch::new(chunk.iter().map(phrase).collect()).expect("a few short names make a batch");
43            let vectors = embedder.embed(&batch).await?;
44            if vectors.len() != chunk.len() {
45                return Err(EmbedError::WrongCount { asked: chunk.len(), got: vectors.len() });
46            }
47            out.extend(chunk.iter().cloned().zip(vectors));
48        }
49        info!("icons: {} picture names embedded", out.len());
50        let _ = self.0.set(out);
51        Ok(self.0.get().expect("just set, or set by another request"))
52    }
53}
54
55/// The `k` icons whose vectors are closest to `query`, closest first. Ties keep the allowlist's order, so the
56/// answer does not depend on how the sort is implemented.
57pub fn shortlist(query: &[f32], index: &[(IconId, Vec<f32>)], k: usize) -> Vec<IconId> {
58    let mut scored: Vec<(f32, usize)> = index.iter().enumerate().map(|(i, (_, v))| (cosine(query, v), i)).collect();
59    scored.sort_by(|a, b| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1)));
60    scored.into_iter().take(k).map(|(_, i)| index[i].0.clone()).collect()
61}
62
63/// Why no picture could be chosen. Not "none fits": that is `Ok(None)`.
64#[derive(Clone, Debug, PartialEq, Eq)]
65pub enum PickError {
66    /// Nothing to choose a picture for (empty), or too long to embed.
67    BadText,
68    Embed(EmbedError),
69    Judge(JudgeError),
70}
71
72impl fmt::Display for PickError {
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        match self {
75            PickError::BadText => f.write_str("the text cannot be given a picture"),
76            PickError::Embed(e) => write!(f, "embedding: {e}"),
77            PickError::Judge(e) => write!(f, "choosing: {e}"),
78        }
79    }
80}
81
82/// The picture for `text`: the best of the shortlist by Jev's judgement, or `None` if none fits.
83pub async fn pick_icon(embedder: &impl Embedder, judge: &impl Judge, vectors: &IconVectors, text: &str) -> Result<Option<IconId>, PickError> {
84    let text = text.trim();
85    if text.is_empty() || text.chars().count() > EMBED_CHARS_MAX {
86        warn!("icons: a text of {} chars cannot be given a picture", text.len());
87        return Err(PickError::BadText);
88    }
89    let index = vectors.get(embedder).await.map_err(PickError::Embed)?;
90    let mut query = embedder.embed(&EmbedBatch::new(vec![text.to_owned()]).map_err(|_| PickError::BadText)?).await.map_err(PickError::Embed)?;
91    let Some(query) = query.pop().filter(|_| query.is_empty()) else {
92        return Err(PickError::Embed(EmbedError::WrongCount { asked: 1, got: query.len() + 1 }));
93    };
94    let candidates = shortlist(&query, index, SHORTLIST);
95    debug!("icons: {} candidates for a text of {} chars", candidates.len(), text.len());
96    let chosen = judge.choose_icon(text, &candidates).await.map_err(PickError::Judge)?;
97    // The contract says the answer is one of the options; an adapter that breaks it gets no picture, not a wrong one.
98    Ok(chosen.filter(|c| candidates.contains(c)))
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    fn id(n: &str) -> IconId {
106        IconId::new(n).unwrap()
107    }
108
109    #[test]
110    fn a_phrase_reads_the_name_as_words() {
111        assert_eq!(phrase(&id("birthday-cake")), "a picture of birthday cake");
112    }
113
114    #[test]
115    fn the_shortlist_is_the_closest_first_and_ties_keep_the_lists_order() {
116        let index = vec![(id("cat"), vec![0.0, 1.0]), (id("dog"), vec![1.0, 0.0]), (id("fox"), vec![1.0, 0.1]), (id("pig"), vec![1.0, 0.0])];
117        let names: Vec<String> = shortlist(&[1.0, 0.0], &index, 3).iter().map(|i| i.to_string()).collect();
118        assert_eq!(names, ["dog", "pig", "fox"]);
119        assert_eq!(shortlist(&[1.0, 0.0], &index, 99).len(), 4);
120        assert!(shortlist(&[1.0, 0.0], &[], 5).is_empty());
121    }
122}