6use serde_json::{Value, json};
7use tokio_postgres::error::SqlState;
8use support::{assert_json_eq, jev_instance};
9use support::mock_jev::{MockJev, Reply};

12const FAVOURED: &str = "skill 200";

15const FIT: f64 = 0.35;

20async fn favouring() -> MockJev {
21    MockJev::start(|req| {
22        let body: Value = serde_json::from_slice(&req.body).unwrap();
23        let mut answers = serde_json::Map::new();
24        for (qid, q) in body["questions"].as_object().unwrap() {
25            if q["type"] == "noul" {
26                answers.insert(qid.clone(), json!({ "type": "noul", "noul": FIT }));
27                continue;
28            }
29            let labels: Vec<&String> = q["criteria"].as_object().unwrap().keys().collect();
30            let n = labels.len() as f64;
31            let favoured = labels.iter().any(|l| *l == FAVOURED);
32            let total = n * (n + 1.0) / 2.0;
33            let probabilities: serde_json::Map<_, _> = labels
34                .iter()
35                .enumerate()
36                .map(|(i, l)| {
37                    let p = match favoured {
38                        true if *l == FAVOURED => 0.9,
39                        true => 0.1 / (n - 1.0),
40                        false => (i as f64 + 1.0) / total,
41                    };
42                    (l.to_string(), json!(p))
43                })
44                .collect();
45            let choice = if favoured { FAVOURED } else { labels[labels.len() - 1] };
46            answers.insert(
47                qid.clone(),
48                json!({ "type": "choice", "choice": choice, "confidence": 0.8, "probabilities": probabilities }),
49            );
50        }
51        Reply::json(200, json!({ "model": "jev-1.13.0", "answers": answers, "usage": { "input_tokens": 300, "output_tokens": 20 } }))
52    })
53    .await
54}

57const OPTIONS: &str = "(SELECT array_agg('skill ' || i ORDER BY i) FROM generate_series(0, 299) i)";
58const DESCRIPTIONS: &str = "(SELECT array_agg('about ' || i ORDER BY i) FROM generate_series(0, 299) i)";

60fn questions(body: &[u8]) -> Vec<Vec<String>> {
61    let body: Value = serde_json::from_slice(body).unwrap();
62    body["questions"]
63        .as_object()
64        .unwrap()
65        .values()
66        .filter(|q| q["type"] == "choice")
67        .map(|q| q["criteria"].as_object().unwrap().keys().cloned().collect())
68        .collect()
69}
70
71#[tokio::test(flavor = "multi_thread")]
72async fn two_requests_then_the_cache() {
73    let mock = favouring().await;
74    let (_pg, client) = jev_instance(&mock).await;
75
76    let sql = format!("SELECT jev_choice_shortlist(t, 'Which skill?', {OPTIONS}, {DESCRIPTIONS}) FROM tickets t");
77    let chosen: String = client.query_one(&sql, &[]).await.unwrap().get(0);
78    assert_eq!(chosen, FAVOURED);
79
80    let requests = mock.requests();
81    assert_eq!(requests.len(), 2, "one request per round");
82
83    // Round 1: both chunks, labels only, the caller's list in order.
84    let chunks = questions(&requests[0].body);
85    assert_eq!(chunks.len(), 2);
86    let sent: Vec<String> = chunks.concat();
87    let caller: Vec<String> = (0..300).map(|i| format!("skill {i}")).collect();
88    assert_eq!(sent, caller);
89    assert!(!String::from_utf8_lossy(&requests[0].body).contains("about "), "chunks send labels only");
90
91    // Round 2: the top 3 of each chunk, in the caller's order, described.
92    let finalists = questions(&requests[1].body);
93    // Chunk 1's ties after the favoured option keep the caller's order.
94    let want: Vec<String> = [147, 148, 149, 150, 151, 200].iter().map(|i| format!("skill {i}")).collect();
95    assert_eq!(finalists, [want]);
96    assert!(String::from_utf8_lossy(&requests[1].body).contains("about 200"), "finalists carry descriptions");
97
98    // One cache row per Choice, and a repeat sends nothing.
99    let cached: i64 = client.query_one("SELECT count(*) FROM jev_cache", &[]).await.unwrap().get(0);
100    assert_eq!(cached, 3);
101    let chosen: String = client.query_one(&sql, &[]).await.unwrap().get(0);
102    assert_eq!(chosen, FAVOURED);
103    assert_eq!(mock.requests().len(), 2, "served from the cache");
104}
105
106#[tokio::test(flavor = "multi_thread")]
107async fn counts_against_max_rows() {
108    let mock = favouring().await;
109    let (_pg, client) = jev_instance(&mock).await;
110    // Estimated at the table's one row, so only the running check refuses.
111    client.batch_execute("ANALYZE tickets; SET jev.max_rows = 1").await.unwrap();
112
113    // The first round is the one row allowed; the second is refused.
114    let err = client
115        .query(&format!("SELECT jev_choice_shortlist(t, 'Which skill?', {OPTIONS}) FROM tickets t"), &[])
116        .await
117        .expect_err("over jev.max_rows");
118    let db = err.as_db_error().expect("an ERROR");
119    assert_eq!(db.code(), &SqlState::PROGRAM_LIMIT_EXCEEDED);
120    assert_eq!(db.message(), "this statement would send more than jev.max_rows (1) rows");
121    assert_eq!(mock.requests().len(), 1);
122}
123
124#[tokio::test(flavor = "multi_thread")]
125async fn a_single_option_is_not_asked() {
126    let mock = favouring().await;
127    let (_pg, client) = jev_instance(&mock).await;
128    let chosen: String = client
129        .query_one("SELECT jev_choice_shortlist(t, 'Which skill?', ARRAY['only']) FROM tickets t", &[])
130        .await
131        .unwrap()
132        .get(0);
133    assert_eq!(chosen, "only");
134    assert!(mock.requests().is_empty());
135}
136
137#[tokio::test(flavor = "multi_thread")]
138async fn bad_options_are_refused_before_sending() {
139    let mock = favouring().await;
140    let (_pg, client) = jev_instance(&mock).await;
141
142    for args in [
143        "ARRAY[]::text[]",
144        "ARRAY['a', NULL]",
145        "ARRAY['a', '']",
146        "ARRAY['a', 'b', 'a']",
147        "ARRAY['a', 'b'], ARRAY['about a']",
148        // 90,000 labels make 355 chunks, whose top 3 are over 255 finalists.
149        "(SELECT array_agg(i::text) FROM generate_series(1, 90000) i)",
150    ] {
151        let err = client
152            .query(&format!("SELECT jev_choice_shortlist(t, 'Which skill?', {args}) FROM tickets t"), &[])
153            .await
154            .expect_err("refused");
155        assert_eq!(err.as_db_error().expect("an ERROR").code(), &SqlState::INVALID_PARAMETER_VALUE, "{args}");
156    }
157    assert!(mock.requests().is_empty(), "refused before sending");
158}

161async fn full(client: &tokio_postgres::Client, args: &str) -> Value {
162    let sql = format!("SELECT to_jsonb(jev_choice_shortlist_full({args}))::text FROM tickets t");
163    let text: String = client.query_one(&sql, &[]).await.unwrap().get(0);
164    serde_json::from_str(&text).unwrap()
165}

167#[tokio::test(flavor = "multi_thread")]
168async fn full_shares_the_cache_rows() {
169    let mock = favouring().await;
170    let (_pg, client) = jev_instance(&mock).await;
171
172    let args = format!("t, 'Which skill?', {OPTIONS}, {DESCRIPTIONS}");
173    let chosen: String =
174        client.query_one(&format!("SELECT jev_choice_shortlist({args}) FROM tickets t"), &[]).await.unwrap().get(0);
175    assert_eq!(chosen, FAVOURED);
176    assert_eq!(mock.requests().len(), 2);
177
178    // The same Choices, so the same three cache rows; the only request
179    // is the "does any label fit" gate, the scalar having asked none.
180    let r = full(&client, &args).await;
181    let requests = mock.requests();
182    assert_eq!(requests.len(), 3, "the Choices served from the cache");
183    assert!(requests[..2].iter().all(|r| nouls(&r.body).is_empty()), "the scalar asked a gate");
184    assert_eq!(nouls(&requests[2].body).len(), 1);
185    assert!(questions(&requests[2].body).is_empty());
186    let cached: i64 = client.query_one("SELECT count(*) FROM jev_cache", &[]).await.unwrap().get(0);
187    assert_eq!(cached, 4);
188    assert_eq!(r["fit"], FIT);
189
190    assert_eq!(r["choice"], FAVOURED);
191    assert_eq!(r["confidence"], 0.8);
192    // The finalists' [label, p] pairs, in the caller's order.
193    let labels: Vec<&str> = r["probabilities"].as_array().unwrap().iter().map(|p| p[0].as_str().unwrap()).collect();
194    assert_eq!(labels, ["skill 147", "skill 148", "skill 149", "skill 150", "skill 151", "skill 200"]);
195    assert_eq!(r["probabilities"][5][1], 0.9);
196    assert_eq!(r["model"], "jev-1.13.0");
197    assert_eq!((r["input_tokens"].as_i64(), r["output_tokens"].as_i64()), (Some(300), Some(20)));
198}
199
200#[tokio::test(flavor = "multi_thread")]
201async fn full_of_a_single_option_is_not_asked() {
202    let mock = favouring().await;
203    let (_pg, client) = jev_instance(&mock).await;
204    let r = full(&client, "t, 'Which skill?', ARRAY['only']").await;
205    assert_json_eq(&r, &json!({
206        "choice": "only", "confidence": 1.0, "probabilities": [["only", 1.0]],
207        "model": null, "input_tokens": 0, "output_tokens": 0, "fit": null,
208    }));
209    assert!(mock.requests().is_empty());
210}

213fn nouls(body: &[u8]) -> Vec<String> {
214    let body: Value = serde_json::from_slice(body).unwrap();
215    body["questions"]
216        .as_object()
217        .unwrap()
218        .values()
219        .filter(|q| q["type"] == "noul")
220        .map(|q| q["instructions"].as_str().unwrap().to_owned())
221        .collect()
222}

224#[tokio::test(flavor = "multi_thread")]
225async fn full_asks_whether_any_label_fits_with_the_chunks() {
226    let mock = favouring().await;
227    let (_pg, client) = jev_instance(&mock).await;
228
229    let r = full(&client, &format!("t, 'Which skill?', {OPTIONS}")).await;
230    assert_eq!(r["choice"], FAVOURED);
231    assert_eq!(r["fit"], FIT);
232
233    // The gate rides in the first request, beside the chunk Choices, and
234    // lists every label in the caller's order; the finalists go alone.
235    let requests = mock.requests();
236    assert_eq!(requests.len(), 2);
237    assert_eq!(questions(&requests[0].body).len(), 2);
238    let every: Vec<String> = (0..300).map(|i| format!("skill {i}")).collect();
239    let want = format!("Which skill?\n\nDoes any of these labels fit? {}", every.join(", "));
240    assert_eq!(nouls(&requests[0].body), [want]);
241    assert!(nouls(&requests[1].body).is_empty());
242}