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

12async fn answering() -> MockJev {
13    MockJev::start(|req| {
14        let body: Value = serde_json::from_slice(&req.body).unwrap();
15        let mut answers = serde_json::Map::new();
16        for (qid, q) in body["questions"].as_object().unwrap() {
17            let answer = match q["type"].as_str() {
18                Some("choice") => {
19                    let labels: Vec<&String> = q["criteria"].as_object().unwrap().keys().collect();
20                    let last = labels.last().unwrap().to_string();
21                    let probabilities: serde_json::Map<_, _> =
22                        labels.iter().map(|l| (l.to_string(), json!(if **l == last { 0.9 } else { 0.05 }))).collect();
23                    json!({ "type": "choice", "choice": last, "confidence": 0.85, "probabilities": probabilities })
24                }
25                Some("score") => {
26                    let n = q["criteria"].as_array().unwrap().len();
27                    let legend: serde_json::Map<_, _> =
28                        (0..n).map(|i| (i.to_string(), q["criteria"][i].clone())).collect();
29                    let probabilities: serde_json::Map<_, _> =
30                        (0..n).map(|i| (i.to_string(), json!(if i == 1 { 1.0 } else { 0.0 }))).collect();
31                    json!({ "type": "score", "score": 1.0, "confidence": 0.8, "legend": legend, "probabilities": probabilities })
32                }
33                _ => json!({ "type": "noul", "noul": 0.75 }),
34            };
35            answers.insert(qid.clone(), answer);
36        }
37        Reply::json(200, json!({ "model": "jev-1.13.0", "answers": answers, "usage": { "input_tokens": 296, "output_tokens": 20 } }))
38    })
39    .await
40}

42const OPTIONS: &str = "ARRAY['sales', 'billing', 'other']";
43
44async fn hits(client: &tokio_postgres::Client) -> i64 {
45    client.query_one("SELECT cache_hits FROM jev_stats()", &[]).await.unwrap().get(0)
46}

49fn record(text: String) -> Value {
50    serde_json::from_str(&text).unwrap()
51}

53#[tokio::test(flavor = "multi_thread")]
54async fn choice_full_shares_jev_choice_judgment_and_cache() {
55    let mock = answering().await;
56    let (_pg, client) = jev_instance(&mock).await;
57
58    let sql = format!(
59        "SELECT jev_choice(t, 'Which team?', {OPTIONS}), to_jsonb(jev_choice_full(t, 'Which team?', {OPTIONS}))::text
60         FROM tickets t"
61    );
62    let row = client.query_one(&sql, &[]).await.unwrap();
63    let expected = json!({
64        "choice": "other",
65        "confidence": 0.85,
66        "probabilities": [["sales", 0.05], ["billing", 0.05], ["other", 0.9]],
67        "model": "jev-1.13.0",
68        "input_tokens": 296,
69        "output_tokens": 20,
70    });
71    assert_eq!(row.get::<_, String>(0), "other");
72    assert_json_eq(&record(row.get(1)), &expected);
73
74    // One request, with the Choice asked once for both calls.
75    let requests = mock.requests();
76    assert_eq!(requests.len(), 1);
77    let sent: Value = serde_json::from_slice(&requests[0].body).unwrap();
78    assert_eq!(sent["questions"].as_object().unwrap().len(), 1, "one question in {sent}");
79
80    // A repeat is served from the cache, with the same fields.
81    let before = hits(&client).await;
82    let row = client.query_one(&sql, &[]).await.unwrap();
83    assert_json_eq(&record(row.get(1)), &expected);
84    assert_eq!(hits(&client).await - before, 2);
85    assert_eq!(mock.requests().len(), 1, "served from the cache");
86}
87
88#[tokio::test(flavor = "multi_thread")]
89async fn eval_and_score_full_fields() {
90    let mock = answering().await;
91    let (_pg, client) = jev_instance(&mock).await;
92
93    let row = client
94        .query_one(
95            "SELECT jev_prob(t, 'Urgent?'), to_jsonb(jev_eval(t, 'Urgent?'))::text,
96                    to_jsonb(jev_score_full(t, 'How angry?', ARRAY['calm', 'annoyed', 'angry']))::text
97             FROM tickets t",
98            &[],
99        )
100        .await
101        .unwrap();
102    assert_eq!(row.get::<_, f64>(0), 0.75);
103    let noul = json!({ "probability": 0.75, "model": "jev-1.13.0", "input_tokens": 296, "output_tokens": 20 });
104    assert_json_eq(&record(row.get(1)), &noul);
105    let score = json!({
106        "score": 1,
107        "confidence": 0.8,
108        "probabilities": [0.0, 1.0, 0.0],
109        "legend": ["calm", "annoyed", "angry"],
110        "model": "jev-1.13.0",
111        "input_tokens": 296,
112        "output_tokens": 20,
113    });
114    assert_json_eq(&record(row.get(2)), &score);
115    // The row is the state, and its two questions went in one request.
116    assert_eq!(mock.requests().len(), 1);
117
118    let before = hits(&client).await;
119    let again: String = client
120        .query_one("SELECT to_jsonb(jev_eval(t, 'Urgent?'))::text FROM tickets t", &[])
121        .await
122        .unwrap()
123        .get(0);
124    assert_json_eq(&record(again), &noul);
125    assert_eq!(hits(&client).await - before, 1);
126    assert_eq!(mock.requests().len(), 1, "served from the cache");
127}
128
129#[tokio::test(flavor = "multi_thread")]
130async fn a_single_option_is_answered_without_a_model() {
131    let mock = answering().await;
132    let (_pg, client) = jev_instance(&mock).await;
133    let text: String = client
134        .query_one("SELECT to_jsonb(jev_choice_full(t, 'Which team?', ARRAY['sales']))::text FROM tickets t", &[])
135        .await
136        .unwrap()
137        .get(0);
138    let expected = json!({
139        "choice": "sales", "confidence": 1.0, "probabilities": [["sales", 1.0]],
140        "model": null, "input_tokens": 0, "output_tokens": 0,
141    });
142    assert_json_eq(&record(text), &expected);
143    assert!(mock.requests().is_empty());
144}