Skip to main content

kunobi_decision/
request.rs

1//! The `systemone` request.
2
3use serde::Serialize;
4use serde_json::{Map, Value};
5
6use crate::error::{Error, Result};
7use crate::questions::Questions;
8use crate::types::Entry;
9
10/// Top-level request fields that extra properties may not shadow.
11const RESERVED_FIELDS: [&str; 3] = ["state", "questions", "model"];
12
13/// State and named questions for [`Client::system_one`](crate::Client::system_one).
14#[derive(Debug, Clone, PartialEq, Serialize)]
15pub struct SystemOneRequest {
16    /// Text, a JSON object or array, or null to evaluate.
17    pub state: Entry,
18    /// Non-empty questions keyed by the names used to identify their answers.
19    pub questions: Questions,
20    /// Model override; `None` uses the client's default model.
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub model: Option<String>,
23    /// Additional top-level properties, forwarded as given.
24    #[serde(flatten)]
25    pub extra: Map<String, Value>,
26}
27
28impl SystemOneRequest {
29    /// A request with the client's default model.
30    pub fn new(state: impl Into<Entry>, questions: Questions) -> Self {
31        Self {
32            state: state.into(),
33            questions,
34            model: None,
35            extra: Map::new(),
36        }
37    }
38
39    /// Use a specific model for this request.
40    pub fn model(mut self, model: impl Into<String>) -> Self {
41        self.model = Some(model.into());
42        self
43    }
44
45    /// Forward an additional top-level property.
46    pub fn extra(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
47        self.extra.insert(key.into(), value.into());
48        self
49    }
50
51    /// Validate the request and serialize it with the model resolved.
52    pub(crate) fn to_body(&self, default_model: &str) -> Result<Vec<u8>> {
53        self.questions.validate()?;
54        if let Some(key) = RESERVED_FIELDS
55            .iter()
56            .find(|field| self.extra.contains_key(**field))
57        {
58            return Err(Error::InvalidRequest(format!(
59                "Extra property \"{key}\" conflicts with a request field; set it directly."
60            )));
61        }
62
63        #[derive(Serialize)]
64        struct Payload<'a> {
65            state: &'a Entry,
66            questions: &'a Questions,
67            model: &'a str,
68            #[serde(flatten)]
69            extra: &'a Map<String, Value>,
70        }
71
72        serde_json::to_vec(&Payload {
73            state: &self.state,
74            questions: &self.questions,
75            model: self.model.as_deref().unwrap_or(default_model),
76            extra: &self.extra,
77        })
78        .map_err(|err| Error::InvalidRequest(format!("Could not serialize the request: {err}")))
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use crate::questions::noul;
86    use serde_json::json;
87
88    fn questions() -> Questions {
89        let mut questions = Questions::new();
90        questions.add("urgent", noul("Urgent?"));
91        questions
92    }
93
94    fn body(request: &SystemOneRequest) -> Value {
95        serde_json::from_slice(&request.to_body("jev-latest").unwrap()).unwrap()
96    }
97
98    #[test]
99    fn resolves_the_default_model() {
100        let request = SystemOneRequest::new("Help ASAP", questions());
101        assert_eq!(
102            body(&request),
103            json!({
104                "state": "Help ASAP",
105                "questions": {"urgent": {"type": "noul", "instructions": "Urgent?"}},
106                "model": "jev-latest"
107            })
108        );
109    }
110
111    #[test]
112    fn model_override_and_extra_properties_are_sent() {
113        let request = SystemOneRequest::new(Entry::Null, questions())
114            .model("jev-2")
115            .extra("seed", 7)
116            .extra("tag", Value::Null);
117        let sent = body(&request);
118        assert_eq!(sent["model"], "jev-2");
119        assert_eq!(sent["seed"], 7);
120        assert_eq!(sent["tag"], Value::Null);
121        assert_eq!(sent["state"], Value::Null);
122    }
123
124    #[test]
125    fn extra_properties_cannot_shadow_request_fields() {
126        for field in RESERVED_FIELDS {
127            let request = SystemOneRequest::new("s", questions()).extra(field, "x");
128            let err = request.to_body("m").unwrap_err();
129            assert!(err.to_string().contains(field), "{err}");
130        }
131    }
132
133    #[test]
134    fn invalid_questions_are_rejected_before_sending() {
135        let request = SystemOneRequest::new("s", Questions::new());
136        assert!(matches!(
137            request.to_body("m"),
138            Err(Error::InvalidRequest(_))
139        ));
140    }
141}