lmjtfy.git / packages / ask / src / tests.rs
tests.rsannotatedtests.rssource335 lines · 13.6 KB · raw
1use std::cell::RefCell;
2use std::collections::VecDeque;
3use std::future::Future;
4use std::time::{Duration, SystemTime};
5
6use bytes::Bytes;
7use jev_client::{Client, HttpRequest, HttpResponse, Instant, Runtime, Transport, TransportError};
8use jev_protocol::ModelId;
9
10use super::*;
11
12/// Replies from a script and keeps what was sent.
13#[derive(Default)]
14struct Wire {
15    replies: RefCell<VecDeque<(u16, String)>>,
16    sent: RefCell<Vec<Vec<u8>>>,
17}
18
19impl Transport for &Wire {
20    async fn send(&self, request: HttpRequest) -> Result<HttpResponse, TransportError> {
21        self.sent.borrow_mut().push(request.body.to_vec());
22        let (status, body) = self.replies.borrow_mut().pop_front().expect("the script ran out");
23        let mut headers = http::HeaderMap::new();
24        headers.insert("x-typesafe-request-id", http::HeaderValue::from_static("req_1"));
25        Ok(HttpResponse { status, headers, body: Bytes::from(body) })
26    }
27
28    fn distrust_connection(&self) {}
29}
30
31#[derive(Clone, Copy)]
32struct Clock(Instant);
33
34impl Runtime for Clock {
35    fn now(&self) -> Instant {
36        self.0
37    }
38    fn wall_clock(&self) -> SystemTime {
39        SystemTime::UNIX_EPOCH
40    }
41    async fn sleep(&self, _: Duration) {}
42    async fn timeout<F: Future>(&self, _: Duration, future: F) -> Option<F::Output> {
43        Some(future.await)
44    }
45    fn jitter(&self) -> f64 {
46        0.0
47    }
48}
49
50fn model() -> ModelId {
51    ModelId::pinned("jev-1.13.0").unwrap()
52}
53
54fn reply(answers: &str) -> String {
55    format!(r#"{{"model":"jev-1.13.0","answers":{answers},"usage":{{"input_tokens":120,"output_tokens":1}}}}"#)
56}
57
58fn block_on<F: Future>(future: F) -> F::Output {
59    let mut future = std::pin::pin!(future);
60    let mut context = std::task::Context::from_waker(std::task::Waker::noop());
61    match future.as_mut().poll(&mut context) {
62        std::task::Poll::Ready(output) => output,
63        std::task::Poll::Pending => panic!("a fake never waits"),
64    }
65}
66
67fn ask(prepared: Prepared, status: u16, body: String) -> (Prepared, Outcome, Vec<u8>) {
68    let wire = Wire::default();
69    wire.replies.borrow_mut().push_back((status, body));
70    let clock = Clock(Instant::now());
71    let client = Client::new(&wire, clock, model(), "key").unwrap();
72    let outcome = block_on(send(&client, &clock, &prepared));
73    let sent = wire.sent.borrow_mut().remove(0);
74    (prepared, outcome, sent)
75}
76
77/// A request for the one fact the old gate was.
78fn gate(input: &str) -> Result<Prepared, ProtocolError> {
79    wanted(&model(), input, &Wanted::Facts(vec![Fact::Answerable]))
80}
81
82fn drafted(input: &str, draft: &Draft) -> Result<Prepared, ProtocolError> {
83    wanted(&model(), input, &Wanted::Drafted(vec![("q1".to_owned(), draft.clone())]))
84}
85
86fn choice() -> Draft {
87    Draft::Choice {
88        instructions: "Which editor is best?".into(),
89        options: vec![
90            llm::Opt { label: "Vim".into(), description: "modal".into() },
91            llm::Opt { label: "Emacs".into(), description: "lisp".into() },
92        ],
93    }
94}
95
96#[test]
97fn the_request_shown_is_the_request_sent() {
98    let (prepared, _, sent) = ask(gate("is water wet?").unwrap(), 200, reply(r#"{"answerable":{"type":"noul","noul":0.9}}"#));
99    assert_eq!(prepared.request.as_bytes(), sent);
100}
101
102#[test]
103fn the_gate_asks_about_the_input_as_state() {
104    let prepared = gate("say \"hi\"").unwrap();
105    let body: serde_json::Value = serde_json::from_str(&prepared.request).unwrap();
106    assert_eq!(body["state"]["input"], "say \"hi\"");
107    assert_eq!(body["model"], "jev-1.13.0");
108    assert_eq!(body["questions"]["answerable"]["type"], "noul");
109    assert!(body["questions"]["answerable"]["criteria"]["true"].is_string());
110    assert!(body["questions"]["answerable"]["criteria"]["false"].is_string());
111    assert_eq!(prepared.parts[0].kind, "noul");
112}
113
114#[test]
115fn a_noul_carries_the_probability_and_the_raw_body() {
116    let raw = reply(r#"{"answerable":{"type":"noul","noul":0.25}}"#);
117    let (_, outcome, _) = ask(gate("is the moon made of cheese?").unwrap(), 200, raw.clone());
118    let Outcome::Answered { judged, body, request_id, attempts, usage, .. } = &outcome else {
119        panic!("{outcome:?}");
120    };
121    assert_eq!(*judged, [Judged::Noul(0.25)]);
122    assert_eq!(*body, raw);
123    assert_eq!(request_id.as_deref(), Some("req_1"));
124    assert_eq!(*attempts, 1);
125    assert_eq!(usage.input_tokens, 120);
126    assert!(outcome.dollars() > 0.0);
127}
128
129#[test]
130fn a_drafted_choice_is_sent_with_its_options_in_order_and_answered() {
131    let prepared = drafted("what is the best editor", &choice()).unwrap();
132    assert_eq!(prepared.parts[0].kind, "choice");
133    let body: serde_json::Value = serde_json::from_str(&prepared.request).unwrap();
134    let criteria = body["questions"]["q1"]["criteria"].as_object().unwrap();
135    assert_eq!(criteria.keys().collect::<Vec<_>>(), ["Vim", "Emacs"]);
136    assert_eq!(criteria["Emacs"], "lisp");
137    let raw = reply(r#"{"q1":{"type":"choice","choice":"Emacs","confidence":0.8,"probabilities":{"Vim":0.1,"Emacs":0.9}}}"#);
138    let (_, outcome, _) = ask(prepared, 200, raw);
139    let Outcome::Answered { judged, .. } = outcome else { panic!("{outcome:?}") };
140    let [Judged::Choice(answer)] = &judged[..] else { panic!("{judged:?}") };
141    assert_eq!(answer.choice, "Emacs");
142    assert_eq!(answer.probabilities, [("Vim".to_owned(), 0.1), ("Emacs".to_owned(), 0.9)]);
143}
144
145#[test]
146fn a_drafted_score_is_answered() {
147    let draft = Draft::Score { instructions: "How good?".into(), levels: vec!["bad".into(), "fine".into(), "good".into()] };
148    let prepared = drafted("rate my idea", &draft).unwrap();
149    assert_eq!(prepared.parts[0].kind, "score");
150    let raw = reply(r#"{"q1":{"type":"score","score":1.4,"confidence":0.6,"legend":{"0":"bad","1":"fine","2":"good"},"probabilities":{"0":0.1,"1":0.4,"2":0.5}}}"#);
151    let (_, outcome, _) = ask(prepared, 200, raw);
152    let Outcome::Answered { judged, .. } = outcome else { panic!("{outcome:?}") };
153    let [Judged::Score(answer)] = &judged[..] else { panic!("{judged:?}") };
154    assert_eq!(answer.score, 1.4);
155    assert_eq!(answer.legend, ["bad", "fine", "good"]);
156}
157
158#[test]
159fn a_single_refusal_is_a_failure_that_cost_nothing() {
160    let (_, outcome, _) = ask(gate("x").unwrap(), 401, r#"{"error":{"message":"bad key"}}"#.into());
161    let Outcome::Failed { error, .. } = &outcome else { panic!("{outcome:?}") };
162    assert!(!error.is_empty());
163    assert_eq!(outcome.dollars(), 0.0);
164}
165
166#[test]
167fn the_worst_case_covers_what_a_call_costs() {
168    let prepared = gate("is water wet?").unwrap();
169    // The live API bills the gate about 350 input tokens (measured 2026-10-02).
170    assert!(prepared.worst_case_dollars >= 350.0 / 1e6 * jev_protocol::DOLLARS_PER_MTOK);
171    assert!(prepared.worst_case_dollars < 0.001);
172}
173
174#[test]
175fn input_is_trimmed_and_capped() {
176    assert_eq!(clean("  hi  "), "hi");
177    assert_eq!(clean(&"é".repeat(MAX_INPUT_CHARS + 10)).chars().count(), MAX_INPUT_CHARS);
178}
179
180#[test]
181fn the_threshold_is_inclusive() {
182    assert!(is_yes(THRESHOLD));
183    assert!(!is_yes(THRESHOLD - 0.01));
184}
185
186#[test]
187fn an_answer_is_sure_at_seventy_percent_either_way() {
188    for sure in [0.0, 0.3, 0.7, 1.0] {
189        assert!(is_sure(sure), "{sure}");
190    }
191    for unsure in [0.31, 0.5, 0.55, 0.69] {
192        assert!(!is_sure(unsure), "{unsure}");
193    }
194}
195
196#[test]
197fn a_kept_response_reads_as_the_same_answer_and_costs_nothing() {
198    let raw = reply(r#"{"answerable":{"type":"noul","noul":0.25}}"#);
199    let (_, sent, _) = ask(gate("is the moon made of cheese?").unwrap(), 200, raw.clone());
200    let Outcome::Answered { judged: first, sent: Sent::Now, .. } = &sent else { panic!("{sent:?}") };
201
202    let prepared = gate("is the moon made of cheese?").unwrap();
203    let kept = Kept {
204        body: &raw,
205        request_id: Some("req_1".into()),
206        attempts: 1,
207        took: Duration::from_millis(150),
208        sent: Sent::Before { at_ms: 1.0 },
209    };
210    let read = read(&model(), &prepared, kept);
211    let Outcome::Answered { judged, body, sent: Sent::Before { .. }, .. } = &read else { panic!("{read:?}") };
212    assert_eq!(judged, first);
213    assert_eq!(*body, raw);
214    assert_eq!(read.dollars(), 0.0);
215}
216
217#[test]
218fn a_kept_response_to_another_question_is_a_failure() {
219    let raw = reply(r#"{"q1":{"type":"noul","noul":0.25}}"#);
220    let prepared = gate("is the moon made of cheese?").unwrap();
221    let kept = Kept { body: &raw, request_id: None, attempts: 1, took: Duration::ZERO, sent: Sent::Before { at_ms: 1.0 } };
222    let read = read(&model(), &prepared, kept);
223    assert!(matches!(read, Outcome::Failed { .. }), "{read:?}");
224}
225
226const READINGS: [Fact; 4] =
227    [Fact::Reads(Kind::Noul), Fact::Reads(Kind::Choice), Fact::Reads(Kind::Score), Fact::Reads(Kind::Chance)];
228
229fn every_jev_fact() -> Vec<Fact> {
230    Fact::ALL.into_iter().filter(|fact| fact.source() == rules::Source::Jev).collect()
231}
232
233#[test]
234fn what_is_wanted_prepares_the_same_request_after_a_round_trip() {
235    for wanted_here in [Wanted::Facts(every_jev_fact()), Wanted::Drafted(vec![("q1".into(), choice())])] {
236        let text = serde_json::to_string(&wanted_here).unwrap();
237        let there: Wanted = serde_json::from_str(&text).unwrap();
238        let input = "what is the best editor";
239        assert_eq!(wanted(&model(), input, &wanted_here).unwrap().request, wanted(&model(), input, &there).unwrap().request);
240    }
241}
242
243#[test]
244fn every_wanted_fact_is_asked_in_one_request_and_the_readings_share_a_question() {
245    let prepared = wanted(&model(), "is water wet?", &Wanted::Facts(every_jev_fact())).unwrap();
246    let body: serde_json::Value = serde_json::from_str(&prepared.request).unwrap();
247    let questions = body["questions"].as_object().unwrap();
248    assert_eq!(
249        questions.keys().collect::<Vec<_>>(),
250        ["answerable", "several", "kind", "scale", "yes", "degree", "chance", "fit"]
251    );
252    assert_eq!(questions["kind"]["type"], "choice");
253    assert_eq!(questions["kind"]["criteria"].as_object().unwrap().len(), 4);
254    assert_eq!(
255        prepared.parts.iter().map(|part| part.kind).collect::<Vec<_>>(),
256        ["noul", "noul", "choice", "noul", "noul", "score", "noul", "noul"]
257    );
258}
259
260fn kinds(yes_or_no: f64, pick_one: f64, how_much: f64) -> Judged {
261    let how_likely = (1.0 - yes_or_no - pick_one - how_much).max(0.0);
262    let probabilities = vec![
263        ("yes_or_no".to_owned(), yes_or_no),
264        ("pick_one".to_owned(), pick_one),
265        ("how_much".to_owned(), how_much),
266        ("how_likely".to_owned(), how_likely),
267    ];
268    let choice = probabilities.iter().max_by(|a, b| a.1.total_cmp(&b.1)).unwrap().0.clone();
269    Judged::Choice(ChoiceAnswer { choice, confidence: 0.5, probabilities })
270}
271
272fn reads(judged: &Judged) -> Vec<bool> {
273    READINGS.iter().map(|fact| learned(*fact, judged) == Some(Value::Bool(true))).collect()
274}
275
276#[test]
277fn a_sure_kind_is_one_reading_and_a_split_is_the_top_two() {
278    assert_eq!(reads(&kinds(0.9, 0.05, 0.05)), [true, false, false, false]);
279    assert_eq!(reads(&kinds(0.4, 0.4, 0.2)), [true, true, false, false]);
280    assert_eq!(reads(&kinds(0.1, 0.6, 0.3)), [false, true, true, false]);
281    // A how-likely question is its own reading.
282    assert_eq!(reads(&kinds(0.05, 0.0, 0.05)), [false, false, false, true]);
283    // The runner-up needs 30%.
284    assert_eq!(reads(&kinds(0.65, 0.06, 0.29)), [true, false, false, false]);
285    // Never three, however even.
286    assert_eq!(reads(&kinds(0.34, 0.33, 0.33)).iter().filter(|read| **read).count(), 2);
287}
288
289#[test]
290fn the_other_facts_are_read_from_their_own_answers() {
291    assert_eq!(learned(Fact::Answerable, &Judged::Noul(0.9)), Some(Value::Bool(true)));
292    assert_eq!(learned(Fact::Several, &Judged::Noul(0.2)), Some(Value::Bool(false)));
293    assert_eq!(learned(Fact::Scale, &Judged::Noul(0.7)), Some(Value::Bool(true)));
294    assert_eq!(learned(Fact::Yes, &Judged::Noul(0.6)), Some(Value::Given));
295    assert_eq!(learned(Fact::Chance, &Judged::Noul(0.28)), Some(Value::Given));
296    assert_eq!(learned(Fact::Reads(Kind::Noul), &Judged::Noul(0.6)), None);
297    assert_eq!(learned(Fact::Answerable, &kinds(1.0, 0.0, 0.0)), None);
298}
299
300#[test]
301fn a_fact_that_is_not_jevs_cannot_be_asked_for() {
302    assert!(wanted(&model(), "x", &Wanted::Facts(vec![Fact::Drafted])).is_err());
303}
304
305#[test]
306fn several_drafts_go_out_as_one_request() {
307    let noul = Draft::Noul { instructions: "Is it?".into(), yes_means: "y".into(), no_means: "n".into() };
308    let prepared = wanted(&model(), "two things", &Wanted::Drafted(vec![("q1".into(), noul), ("q2".into(), choice())])).unwrap();
309    let body: serde_json::Value = serde_json::from_str(&prepared.request).unwrap();
310    assert_eq!(body["questions"].as_object().unwrap().keys().collect::<Vec<_>>(), ["q1", "q2"]);
311    assert_eq!(prepared.id, "answers");
312}
313
314#[test]
315fn a_degree_is_a_score_on_the_fixed_scale_and_fit_is_a_noul() {
316    let prepared = wanted(&model(), "how spicy is a jalapeño", &Wanted::Facts(vec![Fact::Degree, Fact::Fit])).unwrap();
317    assert_eq!(prepared.parts.iter().map(|part| part.kind).collect::<Vec<_>>(), ["score", "noul"]);
318    let raw = reply(
319        r#"{"degree":{"type":"score","score":2.2,"confidence":0.7,"legend":{"0":"Not at all","1":"Slightly","2":"Moderately","3":"Very","4":"Extremely"},"probabilities":{"0":0.0,"1":0.1,"2":0.6,"3":0.3,"4":0.0}},"fit":{"type":"noul","noul":0.99}}"#,
320    );
321    let (_, outcome, _) = ask(prepared, 200, raw);
322    let Outcome::Answered { judged, .. } = outcome else { panic!("{outcome:?}") };
323    assert_eq!(learned(Fact::Degree, &judged[0]), Some(Value::Given));
324    assert_eq!(learned(Fact::Fit, &judged[1]), Some(Value::Bool(true)));
325    let Judged::Score(answer) = &judged[0] else { panic!() };
326    assert_eq!(answer.legend, DEGREES);
327}
328
329#[test]
330fn a_probability_reads_as_a_word_or_two() {
331    for (p, word) in [(0.02, "very unlikely"), (0.28, "unlikely"), (0.5, "a toss-up"), (0.75, "likely"), (0.97, "very likely")] {
332        assert_eq!(likelihood(p), word, "{p}");
333    }
334}
335