1use indexmap::IndexMap;
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Clone, Deserialize, Serialize)]
10pub struct SystemOneResponse {
11 pub model: String,
13 pub answers: IndexMap<String, Answer>,
15 pub usage: Usage,
17}
18
19impl SystemOneResponse {
20 pub fn get(&self, id: &str) -> Option<&Answer> {
22 self.answers.get(id)
23 }
24
25 pub fn nouls(&self) -> impl Iterator<Item = (&str, &NoulAnswer)> {
27 self.answers.iter().filter_map(|(k, v)| match v {
28 Answer::Noul(a) => Some((k.as_str(), a)),
29 _ => None,
30 })
31 }
32
33 pub fn choices(&self) -> impl Iterator<Item = (&str, &ChoiceAnswer)> {
35 self.answers.iter().filter_map(|(k, v)| match v {
36 Answer::Choice(a) => Some((k.as_str(), a)),
37 _ => None,
38 })
39 }
40
41 pub fn scores(&self) -> impl Iterator<Item = (&str, &ScoreAnswer)> {
43 self.answers.iter().filter_map(|(k, v)| match v {
44 Answer::Score(a) => Some((k.as_str(), a)),
45 _ => None,
46 })
47 }
48
49 pub fn noul(&self, id: &str) -> Option<&NoulAnswer> {
51 self.answers.get(id).and_then(Answer::as_noul)
52 }
53
54 pub fn choice(&self, id: &str) -> Option<&ChoiceAnswer> {
56 self.answers.get(id).and_then(Answer::as_choice)
57 }
58
59 pub fn score(&self, id: &str) -> Option<&ScoreAnswer> {
61 self.answers.get(id).and_then(Answer::as_score)
62 }
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
67pub struct Usage {
68 pub input_tokens: u64,
70 pub output_tokens: u64,
72}
73
74#[derive(Debug, Clone, Deserialize, Serialize)]
76#[serde(tag = "type", rename_all = "lowercase")]
77pub enum Answer {
78 Noul(NoulAnswer),
80 Choice(ChoiceAnswer),
82 Score(ScoreAnswer),
84}
85
86impl Answer {
87 pub fn as_noul(&self) -> Option<&NoulAnswer> {
89 match self {
90 Answer::Noul(a) => Some(a),
91 _ => None,
92 }
93 }
94
95 pub fn as_choice(&self) -> Option<&ChoiceAnswer> {
97 match self {
98 Answer::Choice(a) => Some(a),
99 _ => None,
100 }
101 }
102
103 pub fn as_score(&self) -> Option<&ScoreAnswer> {
105 match self {
106 Answer::Score(a) => Some(a),
107 _ => None,
108 }
109 }
110}
111
112#[derive(Debug, Clone, Copy, Deserialize, Serialize)]
114pub struct NoulAnswer {
115 pub noul: f64,
117}
118
119#[derive(Debug, Clone, Deserialize, Serialize)]
121pub struct ChoiceAnswer {
122 pub choice: String,
124 #[serde(default)]
126 pub probabilities: IndexMap<String, f64>,
127 pub confidence: f64,
129}
130
131#[derive(Debug, Clone, Deserialize, Serialize)]
133pub struct ScoreAnswer {
134 pub score: f64,
137 #[serde(default)]
139 pub legend: IndexMap<String, String>,
140 #[serde(default)]
142 pub probabilities: IndexMap<String, f64>,
143 pub confidence: f64,
145}
146
147#[cfg(test)]
148mod tests {
149 use super::*;
150 use serde_json::json;
151
152 #[test]
153 fn parses_full_response() {
154 let raw = json!({
155 "model": "jev-1.13.0",
156 "answers": {
157 "department": {
158 "type": "choice",
159 "choice": "technical",
160 "probabilities": {"billing": 0.159, "technical": 0.84, "sales": 0.001},
161 "confidence": 0.596,
162 },
163 "frustration": {
164 "type": "score",
165 "score": 1.035,
166 "legend": {"0": "Calm", "1": "Frustrated", "2": "Very angry"},
167 "probabilities": {"0": 0.1, "1": 0.6, "2": 0.3},
168 "confidence": 0.842,
169 },
170 "is_urgent": {"type": "noul", "noul": 0.999},
171 },
172 "usage": {"input_tokens": 312, "output_tokens": 48},
173 });
174
175 let resp: SystemOneResponse = serde_json::from_value(raw).unwrap();
176 assert_eq!(resp.model, "jev-1.13.0");
177 assert_eq!(resp.answers.len(), 3);
178 assert_eq!(resp.usage.input_tokens, 312);
179
180 let dept = resp.choice("department").unwrap();
181 assert_eq!(dept.choice, "technical");
182 assert_eq!(dept.probabilities["billing"], 0.159);
183 assert_eq!(dept.confidence, 0.596);
184
185 let fr = resp.score("frustration").unwrap();
186 assert_eq!(fr.score, 1.035);
187 assert_eq!(fr.legend["2"], "Very angry");
188 assert_eq!(fr.probabilities["1"], 0.6);
189
190 let urg = resp.noul("is_urgent").unwrap();
191 assert_eq!(urg.noul, 0.999);
192
193 assert_eq!(resp.nouls().count(), 1);
194 assert_eq!(resp.choices().count(), 1);
195 assert_eq!(resp.scores().count(), 1);
196 }
197
198 #[test]
199 fn score_answer_tolerates_missing_probabilities() {
200 let raw = json!({
201 "model": "jev-latest",
202 "answers": {
203 "s": {"type": "score", "score": 0.0, "confidence": 1.0, "legend": {"0": "only"}}
204 },
205 "usage": {"input_tokens": 1, "output_tokens": 1},
206 });
207 let resp: SystemOneResponse = serde_json::from_value(raw).unwrap();
208 assert!(resp.score("s").unwrap().probabilities.is_empty());
209 }
210
211 #[test]
212 fn responses_round_trip_through_json() {
213 let raw = json!({
215 "model": "jev-1.13.0",
216 "answers": {
217 "ok": {"type": "noul", "noul": 0.87},
218 "path": {
219 "type": "choice",
220 "choice": "a",
221 "probabilities": {"a": 0.9, "b": 0.1},
222 "confidence": 0.8,
223 },
224 },
225 "usage": {"input_tokens": 10, "output_tokens": 4},
226 });
227 let resp: SystemOneResponse = serde_json::from_value(raw.clone()).unwrap();
228 assert_eq!(serde_json::to_value(&resp).unwrap(), raw);
229 }
230}