forge_system_one/
wire.rs

1//! The System One wire shapes: one typed question out, one typed answer
2//! back, and the check an answer must survive before it is trusted.
3
4use std::collections::{BTreeMap, BTreeSet};
5
6/// One typed question; serialized as the request's `questions` value.
7/// `instructions` and the criteria values the API types as any JSON
8/// (string, object, array or null; a score level no null) cross verbatim.
9#[derive(Debug, Clone, PartialEq, serde::Serialize)]
10#[serde(tag = "type", rename_all = "lowercase")]
11pub enum Question {
12    Noul {
13        instructions: serde_json::Value,
14        #[serde(skip_serializing_if = "Option::is_none")]
15        criteria: Option<NoulCriteria>,
16    },
17    Choice {
18        instructions: serde_json::Value,
19        criteria: BTreeMap<String, serde_json::Value>,
20    },
21    Score {
22        instructions: serde_json::Value,
23        criteria: Vec<serde_json::Value>,
24    },
25}
26
27/// What yes and no mean for a noul question, when the boundary matters;
28/// the keys are fixed, the values are any JSON.
29#[derive(Debug, Clone, PartialEq, serde::Serialize)]
30pub struct NoulCriteria {
31    #[serde(rename = "true")]
32    pub r#true: serde_json::Value,
33    #[serde(rename = "false")]
34    pub r#false: serde_json::Value,
35}
36
37/// One typed answer, deserialized from the response.
38#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
39#[serde(tag = "type", rename_all = "lowercase")]
40pub enum Answer {
41    Noul {
42        noul: f64,
43    },
44    Choice {
45        choice: String,
46        #[serde(skip_serializing_if = "Option::is_none")]
47        probabilities: Option<BTreeMap<String, f64>>,
48        #[serde(skip_serializing_if = "Option::is_none")]
49        confidence: Option<f64>,
50    },
51    Score {
52        score: f64,
53        #[serde(skip_serializing_if = "Option::is_none")]
54        probabilities: Option<BTreeMap<String, f64>>,
55        #[serde(skip_serializing_if = "Option::is_none")]
56        confidence: Option<f64>,
57        #[serde(skip_serializing_if = "Option::is_none")]
58        legend: Option<BTreeMap<String, String>>,
59    },
60}
61
62/// Token usage as the provider reports it; `cost` only on OpenRouter.
63#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
64pub struct Usage {
65    pub input_tokens: u64,
66    pub output_tokens: u64,
67    #[serde(skip_serializing_if = "Option::is_none")]
68    pub cost: Option<f64>,
69}
70
71/// A resolved decision: what answered, the answer, and what it cost, if
72/// the provider reported a usage block.
73#[derive(Debug, Clone, PartialEq)]
74pub struct AskOutcome {
75    pub model: String,
76    pub answer: Answer,
77    pub usage: Option<Usage>,
78}
79
80/// The answer must be internally consistent with the question it answers:
81/// keys cover the criteria, probabilities are probabilities, the score is
82/// inside its levels. Anything else is a provider bug, not a decision.
83pub fn validate_answer(question: &Question, answer: &Answer) -> Result<(), String> {
84    match (question, answer) {
85        (Question::Noul { .. }, Answer::Noul { noul }) => check_unit(*noul, "noul"),
86        (
87            Question::Choice { criteria, .. },
88            Answer::Choice { choice, probabilities, confidence },
89        ) => {
90            if !criteria.contains_key(choice) {
91                return Err(format!("choice `{choice}` is not one of the criteria keys"));
92            }
93            if let Some(probabilities) = probabilities {
94                let expected: BTreeSet<String> = criteria.keys().cloned().collect();
95                check_probability_map(probabilities, &expected)?;
96            }
97            check_confidence(*confidence)
98        }
99        (
100            Question::Score { criteria, .. },
101            Answer::Score { score, probabilities, confidence, legend },
102        ) => {
103            let levels = criteria.len();
104            let Some(last) = levels.checked_sub(1) else {
105                return Err("score question defines no levels".to_owned());
106            };
107            let Ok(last) = u32::try_from(last) else {
108                return Err("score question has more levels than supported".to_owned());
109            };
110            let score = *score;
111            if !score.is_finite() || score < 0.0 || score > f64::from(last) {
112                return Err(format!("score {score} outside 0..={last}"));
113            }
114            let expected: BTreeSet<String> = (0..=last).map(|level| level.to_string()).collect();
115            if let Some(probabilities) = probabilities {
116                check_probability_map(probabilities, &expected)?;
117            }
118            if let Some(legend) = legend {
119                let keys: BTreeSet<&String> = legend.keys().collect();
120                let want: BTreeSet<&String> = expected.iter().collect();
121                if keys != want {
122                    return Err(format!("legend keys {keys:?} do not match the levels {want:?}"));
123                }
124            }
125            check_confidence(*confidence)
126        }
127        _ => Err("answer type does not match the question type".to_owned()),
128    }
129}
130
131/// A 0..=1 field: noul, a probability, a confidence.
132fn check_unit(value: f64, field: &str) -> Result<(), String> {
133    if !value.is_finite() || !(0.0..=1.0).contains(&value) {
134        return Err(format!("{field} {value} outside 0..=1"));
135    }
136    Ok(())
137}
138
139fn check_confidence(confidence: Option<f64>) -> Result<(), String> {
140    match confidence {
141        Some(confidence) => check_unit(confidence, "confidence"),
142        None => Ok(()),
143    }
144}
145
146/// Keys, values, and the sum tolerance pydantic-ai applies for the
147/// provider's two-decimal display: 1e-6 + n * 0.005.
148fn check_probability_map(
149    probabilities: &BTreeMap<String, f64>,
150    expected: &BTreeSet<String>,
151) -> Result<(), String> {
152    let keys: BTreeSet<&String> = probabilities.keys().collect();
153    let want: BTreeSet<&String> = expected.iter().collect();
154    if keys != want {
155        return Err(format!("probability keys {keys:?} do not match the expected keys {want:?}"));
156    }
157    let mut sum = 0.0;
158    for value in probabilities.values() {
159        check_unit(*value, "probability")?;
160        sum += value;
161    }
162    let count = u32::try_from(probabilities.len()).unwrap_or(u32::MAX);
163    let tolerance = 1e-6 + f64::from(count) * 0.005;
164    if (sum - 1.0).abs() > tolerance {
165        return Err(format!("probabilities sum to {sum} (expected ~1)"));
166    }
167    Ok(())
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    fn noul_question() -> Question {
175        Question::Noul { instructions: serde_json::json!("Is this a bug?"), criteria: None }
176    }
177
178    fn choice_question() -> Question {
179        Question::Choice {
180            instructions: serde_json::json!("Which team?"),
181            criteria: [
182                ("billing".to_owned(), serde_json::json!("Payments and refunds")),
183                ("technical".to_owned(), serde_json::Value::Null),
184            ]
185            .into_iter()
186            .collect(),
187        }
188    }
189
190    fn score_question() -> Question {
191        Question::Score {
192            instructions: serde_json::json!("How urgent?"),
193            criteria: vec![
194                serde_json::json!("Routine"),
195                serde_json::json!("Soon"),
196                serde_json::json!("Urgent"),
197            ],
198        }
199    }
200
201    fn probs(pairs: &[(&str, f64)]) -> std::collections::BTreeMap<String, f64> {
202        pairs.iter().map(|(key, value)| ((*key).to_owned(), *value)).collect()
203    }
204
205    #[test]
206    fn question_serializes_to_the_wire_shape() {
207        let q = Question::Noul {
208            instructions: serde_json::json!("Is this a billing issue?"),
209            criteria: Some(NoulCriteria {
210                r#true: serde_json::json!("Payments or refunds"),
211                r#false: serde_json::json!("Anything else"),
212            }),
213        };
214        assert_eq!(
215            serde_json::to_value(&q).unwrap(),
216            serde_json::json!({"type":"noul","instructions":"Is this a billing issue?","criteria":{"true":"Payments or refunds","false":"Anything else"}})
217        );
218        let c = Question::Choice {
219            instructions: serde_json::json!("Which team?"),
220            criteria: [
221                ("billing".to_owned(), serde_json::json!("Payments")),
222                ("frontend".to_owned(), serde_json::Value::Null),
223            ]
224            .into_iter()
225            .collect(),
226        };
227        assert_eq!(
228            serde_json::to_value(&c).unwrap(),
229            serde_json::json!({"type":"choice","instructions":"Which team?","criteria":{"billing":"Payments","frontend":null}})
230        );
231        let s = Question::Score {
232            instructions: serde_json::json!("How urgent?"),
233            criteria: vec![
234                serde_json::json!("Routine"),
235                serde_json::json!("Soon"),
236                serde_json::json!("Urgent"),
237            ],
238        };
239        assert_eq!(
240            serde_json::to_value(&s).unwrap(),
241            serde_json::json!({"type":"score","instructions":"How urgent?","criteria":["Routine","Soon","Urgent"]})
242        );
243    }
244
245    /// The API types `instructions` and criteria values as any JSON; the
246    /// structure must survive serialization untouched.
247    #[test]
248    fn structured_values_serialize_to_the_wire_shape() {
249        let q = Question::Choice {
250            instructions: serde_json::json!({"question": "Which team?", "touched": ["a.rs"]}),
251            criteria: [
252                ("billing".to_owned(), serde_json::json!({"files": ["a.rs"]})),
253                ("frontend".to_owned(), serde_json::Value::Null),
254            ]
255            .into_iter()
256            .collect(),
257        };
258        assert_eq!(
259            serde_json::to_value(&q).unwrap(),
260            serde_json::json!({
261                "type": "choice",
262                "instructions": {"question": "Which team?", "touched": ["a.rs"]},
263                "criteria": {"billing": {"files": ["a.rs"]}, "frontend": null}
264            })
265        );
266
267        let s = Question::Score {
268            instructions: serde_json::json!({"question": "How urgent?", "note": "per `sev.md`"}),
269            criteria: vec![
270                serde_json::json!("Routine"),
271                serde_json::json!({"label": "Urgent", "scope": "page now"}),
272            ],
273        };
274        assert_eq!(
275            serde_json::to_value(&s).unwrap(),
276            serde_json::json!({
277                "type": "score",
278                "instructions": {"question": "How urgent?", "note": "per `sev.md`"},
279                "criteria": ["Routine", {"label": "Urgent", "scope": "page now"}]
280            })
281        );
282    }
283
284    #[test]
285    fn noul_without_criteria_omits_the_field() {
286        assert_eq!(
287            serde_json::to_value(noul_question()).unwrap(),
288            serde_json::json!({"type":"noul","instructions":"Is this a bug?"})
289        );
290    }
291
292    #[test]
293    fn noul_answer_ignores_extra_payload_fields() {
294        let answer: Answer =
295            serde_json::from_value(serde_json::json!({"type":"noul","noul":0.5,"extra":1}))
296                .unwrap();
297        assert_eq!(answer, Answer::Noul { noul: 0.5 });
298    }
299
300    #[test]
301    fn choice_answer_probabilities_default_to_none() {
302        let answer: Answer =
303            serde_json::from_value(serde_json::json!({"type":"choice","choice":"billing"}))
304                .unwrap();
305        assert_eq!(
306            answer,
307            Answer::Choice { choice: "billing".to_owned(), probabilities: None, confidence: None }
308        );
309    }
310
311    /// The tool output serializes an answer back out; absent optionals
312    /// must not come back as nulls.
313    #[test]
314    fn answer_serializes_without_absent_option_fields() {
315        let answer =
316            Answer::Choice { choice: "billing".to_owned(), probabilities: None, confidence: None };
317        assert_eq!(
318            serde_json::to_value(&answer).unwrap(),
319            serde_json::json!({"type":"choice","choice":"billing"})
320        );
321        let usage = Usage { input_tokens: 10, output_tokens: 3, cost: None };
322        assert_eq!(
323            serde_json::to_value(&usage).unwrap(),
324            serde_json::json!({"input_tokens":10,"output_tokens":3})
325        );
326    }
327
328    #[test]
329    fn noul_in_range_validates() {
330        assert_eq!(validate_answer(&noul_question(), &Answer::Noul { noul: 0.83 }), Ok(()));
331    }
332
333    #[test]
334    fn noul_out_of_range_is_refused() {
335        assert_eq!(
336            validate_answer(&noul_question(), &Answer::Noul { noul: 1.5 }),
337            Err("noul 1.5 outside 0..=1".to_owned())
338        );
339    }
340
341    #[test]
342    fn noul_nan_is_refused() {
343        assert_eq!(
344            validate_answer(&noul_question(), &Answer::Noul { noul: f64::NAN }),
345            Err("noul NaN outside 0..=1".to_owned())
346        );
347    }
348
349    #[test]
350    fn choice_in_criteria_validates() {
351        let answer = Answer::Choice {
352            choice: "billing".to_owned(),
353            probabilities: Some(probs(&[("billing", 0.9), ("technical", 0.1)])),
354            confidence: Some(0.8),
355        };
356        assert_eq!(validate_answer(&choice_question(), &answer), Ok(()));
357    }
358
359    #[test]
360    fn choice_outside_criteria_is_refused() {
361        let answer =
362            Answer::Choice { choice: "nope".to_owned(), probabilities: None, confidence: None };
363        assert_eq!(
364            validate_answer(&choice_question(), &answer),
365            Err("choice `nope` is not one of the criteria keys".to_owned())
366        );
367    }
368
369    #[test]
370    fn choice_probability_key_drift_is_refused() {
371        let answer = Answer::Choice {
372            choice: "billing".to_owned(),
373            probabilities: Some(probs(&[("billing", 0.9), ("sales", 0.1)])),
374            confidence: None,
375        };
376        let err =
377            validate_answer(&choice_question(), &answer).expect_err("drifted keys are refused");
378        assert!(err.starts_with("probability keys"), "{err}");
379    }
380
381    #[test]
382    fn probabilities_wrong_sum_is_refused() {
383        let answer = Answer::Choice {
384            choice: "billing".to_owned(),
385            probabilities: Some(probs(&[("billing", 0.5), ("technical", 0.2)])),
386            confidence: None,
387        };
388        assert_eq!(
389            validate_answer(&choice_question(), &answer),
390            Err("probabilities sum to 0.7 (expected ~1)".to_owned())
391        );
392    }
393
394    #[test]
395    fn probabilities_within_display_tolerance_validate() {
396        // 0.005 of drift on two options sits inside 1e-6 + 2 * 0.005.
397        let answer = Answer::Choice {
398            choice: "billing".to_owned(),
399            probabilities: Some(probs(&[("billing", 0.995), ("technical", 0.01)])),
400            confidence: None,
401        };
402        assert_eq!(validate_answer(&choice_question(), &answer), Ok(()));
403    }
404
405    #[test]
406    fn probability_value_out_of_range_is_refused() {
407        let answer = Answer::Choice {
408            choice: "billing".to_owned(),
409            probabilities: Some(probs(&[("billing", 1.5), ("technical", -0.5)])),
410            confidence: None,
411        };
412        let err = validate_answer(&choice_question(), &answer)
413            .expect_err("out-of-range probability is refused");
414        assert!(err.starts_with("probability 1.5 outside 0..=1"), "{err}");
415    }
416
417    #[test]
418    fn confidence_out_of_range_is_refused() {
419        let answer = Answer::Choice {
420            choice: "billing".to_owned(),
421            probabilities: None,
422            confidence: Some(1.2),
423        };
424        assert_eq!(
425            validate_answer(&choice_question(), &answer),
426            Err("confidence 1.2 outside 0..=1".to_owned())
427        );
428    }
429
430    #[test]
431    fn score_in_range_validates() {
432        let answer = Answer::Score {
433            score: 1.79,
434            probabilities: Some(probs(&[("0", 0.1), ("1", 0.3), ("2", 0.6)])),
435            confidence: Some(0.5),
436            legend: Some(
437                [
438                    ("0".to_owned(), "Routine".to_owned()),
439                    ("1".to_owned(), "Soon".to_owned()),
440                    ("2".to_owned(), "Urgent".to_owned()),
441                ]
442                .into_iter()
443                .collect(),
444            ),
445        };
446        assert_eq!(validate_answer(&score_question(), &answer), Ok(()));
447    }
448
449    #[test]
450    fn score_outside_levels_is_refused() {
451        let answer =
452            Answer::Score { score: 2.5, probabilities: None, confidence: None, legend: None };
453        assert_eq!(
454            validate_answer(&score_question(), &answer),
455            Err("score 2.5 outside 0..=2".to_owned())
456        );
457    }
458
459    #[test]
460    fn score_probability_key_drift_is_refused() {
461        let answer = Answer::Score {
462            score: 1.0,
463            probabilities: Some(probs(&[("0", 0.5), ("2", 0.5)])),
464            confidence: None,
465            legend: None,
466        };
467        let err = validate_answer(&score_question(), &answer)
468            .expect_err("a missing level key is refused");
469        assert!(err.starts_with("probability keys"), "{err}");
470    }
471
472    #[test]
473    fn legend_key_drift_is_refused() {
474        let answer = Answer::Score {
475            score: 1.0,
476            probabilities: None,
477            confidence: None,
478            legend: Some(
479                [("0".to_owned(), "Routine".to_owned()), ("5".to_owned(), "Nope".to_owned())]
480                    .into_iter()
481                    .collect(),
482            ),
483        };
484        let err =
485            validate_answer(&score_question(), &answer).expect_err("a drifted legend is refused");
486        assert!(err.starts_with("legend keys"), "{err}");
487    }
488
489    #[test]
490    fn answer_type_mismatch_is_refused() {
491        let answer =
492            Answer::Choice { choice: "billing".to_owned(), probabilities: None, confidence: None };
493        assert_eq!(
494            validate_answer(&noul_question(), &answer),
495            Err("answer type does not match the question type".to_owned())
496        );
497    }
498
499    #[test]
500    fn score_question_with_no_levels_is_refused() {
501        let question = Question::Score {
502            instructions: serde_json::json!("How urgent?"),
503            criteria: Vec::new(),
504        };
505        let answer =
506            Answer::Score { score: 0.0, probabilities: None, confidence: None, legend: None };
507        assert_eq!(
508            validate_answer(&question, &answer),
509            Err("score question defines no levels".to_owned())
510        );
511    }
512}
forge_system_one/wire.rs

forge_system_one/
wire.rs
