lmjtfy.git / packages / llm / src / tests.rs
tests.rsannotatedtests.rssource151 lines · 6.5 KB · raw
1use super::*;
2
3fn reply(tool_calls: &str) -> String {
4    format!(
5        r#"{{"choices":[{{"index":0,"message":{{"role":"assistant","content":null,"tool_calls":{tool_calls}}}}}],"usage":{{"prompt_tokens":700,"completion_tokens":150,"neurons":12.5}}}}"#
6    )
7}
8
9fn call(name: &str, arguments: &str) -> String {
10    serde_json::json!([{ "id": "c1", "type": "function", "function": { "name": name, "arguments": arguments } }])
11        .to_string()
12}
13
14fn only(body: &str) -> ToolCall {
15    let mut reply = parse(body).unwrap();
16    assert_eq!(reply.calls.len(), 1);
17    reply.calls.remove(0)
18}
19
20#[test]
21fn the_request_carries_the_input_and_all_three_tools() {
22    let body: Value = serde_json::from_str(&request("is \"water\" wet?", &[])).unwrap();
23    assert_eq!(body["messages"][1]["content"], "is \"water\" wet?");
24    let names: Vec<_> = body["tools"].as_array().unwrap().iter().map(|t| t["function"]["name"].clone()).collect();
25    assert_eq!(names, ["jev_noul", "jev_choice", "jev_score"]);
26    assert_eq!(body["max_tokens"], MAX_TOKENS);
27}
28
29#[test]
30fn a_choice_is_read() {
31    let made = only(&reply(&call(
32        "jev_choice",
33        r#"{"instructions":"Which editor is best?","options":[{"label":"Vim","description":"modal"},{"label":"Emacs","description":"lisp"}]}"#,
34    )));
35    let Ok(Draft::Choice { instructions, options }) = made.draft else { panic!("{made:?}") };
36    assert_eq!(instructions, "Which editor is best?");
37    assert_eq!(options.len(), 2);
38    assert_eq!(options[1].label, "Emacs");
39}
40
41#[test]
42fn rest_wraps_the_reply_in_result() {
43    let wrapped = format!(r#"{{"result":{},"success":true}}"#, reply(&call("jev_score", r#"{"instructions":"How good?","levels":["bad","fine","good"]}"#)));
44    let parsed = parse(&wrapped).unwrap();
45    assert!(matches!(parsed.calls[0].draft, Ok(Draft::Score { .. })));
46    assert_eq!(parsed.usage, Usage { prompt_tokens: 700, completion_tokens: 150, neurons: Some(12.5) });
47}
48
49#[test]
50fn arguments_encoded_twice_are_decoded_once_more() {
51    let once = r#"{"instructions":"Is it?","yes_means":"it is","no_means":"it is not"}"#;
52    let twice = serde_json::to_string(once).unwrap();
53    let made = only(&reply(&call("jev_noul", &twice)));
54    assert!(matches!(made.draft, Ok(Draft::Noul { .. })), "{made:?}");
55    assert_eq!(made.arguments, twice, "the arguments are kept as the model wrote them");
56}
57
58#[test]
59fn malformed_calls_say_why() {
60    for (name, arguments, why) in [
61        ("jev_choice", r#"{"instructions":"q","options":[{"label":"only","description":"one"}]}"#, "2 to 8 options"),
62        ("jev_choice", r#"{"instructions":"q","options":[{"label":"a","description":""},{"label":"a","description":""}]}"#, "appears twice"),
63        ("jev_score", r#"{"instructions":"q","levels":["one"]}"#, "2 to 10 levels"),
64        ("jev_noul", r#"{"instructions":"","yes_means":"y","no_means":"n"}"#, "instructions is empty"),
65        ("jev_noul", r#"{"instructions":"q","yes_means":"y"}"#, "no_means"),
66        ("jev_noul", r#"{"instructions":"q","yes_means":"y","no_means":"n","extra":1}"#, "unknown field"),
67        ("jev_noul", "", "not JSON"),
68        ("jev_noul", "[1]", "not an object"),
69        ("jev_write", r#"{"text":"hi"}"#, "no tool called"),
70    ] {
71        let made = only(&reply(&call(name, arguments)));
72        let error = made.draft.expect_err(arguments);
73        assert!(error.contains(why), "{arguments}: {error}");
74    }
75}
76
77#[test]
78fn a_reply_with_no_tool_call_has_no_calls() {
79    let parsed = parse(r#"{"choices":[{"message":{"role":"assistant","content":"Vim, obviously."}}]}"#).unwrap();
80    assert!(parsed.calls.is_empty());
81    assert_eq!(parsed.usage, Usage::default());
82}
83
84#[test]
85fn calls_past_the_cap_are_dropped_and_counted() {
86    let one = serde_json::json!({ "type": "function", "function": { "name": "jev_noul", "arguments": r#"{"instructions":"q","yes_means":"y","no_means":"n"}"# } });
87    let six = Value::Array(vec![one; 6]).to_string();
88    let parsed = parse(&reply(&six)).unwrap();
89    assert_eq!(parsed.calls.len(), MAX_CALLS);
90    assert_eq!(parsed.dropped, 2);
91}
92
93#[test]
94fn a_shapeless_reply_is_an_error() {
95    assert!(parse("upstream down").is_err());
96    assert!(parse(r#"{"errors":[{"message":"no"}]}"#).is_err());
97}
98
99#[test]
100fn neurons_come_from_the_reply_or_the_price_list() {
101    let model = Model::find("@cf/ibm-granite/granite-4.0-h-micro").unwrap();
102    let reported = Usage { prompt_tokens: 1_000_000, completion_tokens: 0, neurons: Some(7.0) };
103    assert_eq!(model.neurons(reported), 7.0);
104    let priced = Usage { prompt_tokens: 1_000_000, completion_tokens: 1_000_000, neurons: None };
105    assert_eq!(model.neurons(priced), 1542.0 + 10158.0);
106    assert!(model.worst_case_neurons(&request("hi", &[])) < 50.0);
107}
108
109fn sent(wants: &[Want]) -> (String, Vec<String>) {
110    let body: Value = serde_json::from_str(&request("x", wants)).unwrap();
111    let system = body["messages"][0]["content"].as_str().unwrap().to_owned();
112    let tools = body["tools"].as_array().unwrap().iter().map(|tool| tool["function"]["name"].as_str().unwrap().to_owned()).collect();
113    (system, tools)
114}
115
116#[test]
117fn what_the_rules_settled_is_all_the_llm_is_told_about_and_given() {
118    let (system, tools) = sent(&[Want::Scale]);
119    assert_eq!(tools, ["jev_score"]);
120    assert!(system.contains("exactly one jev_score call"));
121    // Nothing invites the yes-or-no Jev has already answered itself.
122    assert!(!system.contains("jev_noul") && !system.contains("jev_choice"), "{system}");
123
124    let (system, tools) = sent(&[Want::Options]);
125    assert_eq!(tools, ["jev_choice"]);
126    assert!(system.contains("exactly one jev_choice call") && !system.contains("jev_noul"));
127
128    let (system, tools) = sent(&[Want::Options, Want::Scale]);
129    assert_eq!(tools, ["jev_choice", "jev_score"]);
130    assert!(system.contains("exactly two tool calls"));
131}
132
133#[test]
134fn several_questions_are_left_to_the_llm_with_every_tool() {
135    let all = sent(&[]);
136    assert_eq!(all.1, ["jev_noul", "jev_choice", "jev_score"]);
137    assert_eq!(sent(&[Want::Split]), all);
138    assert_eq!(sent(&[Want::Options, Want::Split]), all);
139}
140
141#[test]
142fn a_question_of_a_kind_that_was_not_asked_for_is_not_taken() {
143    let noul = Draft::Noul { instructions: "Is it?".into(), yes_means: "y".into(), no_means: "n".into() };
144    let score = Draft::Score { instructions: "How?".into(), levels: vec!["low".into(), "high".into()] };
145    assert!(!takes(&[Want::Scale], &noul));
146    assert!(takes(&[Want::Scale], &score));
147    assert!(!takes(&[Want::Options], &score));
148    assert!(takes(&[Want::Options, Want::Scale], &score));
149    assert!(takes(&[Want::Split], &noul));
150    assert!(takes(&[], &noul));
151}