1use crate::{Answer, AnswerKind, Error, MultiObjective, Question, QuestionSpec, QuestionType};
3use indexmap::IndexMap;
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7pub type Automation = IndexMap<String, Option<bool>>;
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum Verdict {
13 Supported,
15 Contradicted,
17 NotStated,
19 NotSupported,
21}
22impl Verdict {
23 fn key(self) -> &'static str {
24 match self {
25 Self::Supported => "supported",
26 Self::Contradicted => "contradicted",
27 Self::NotStated => "not_stated",
28 Self::NotSupported => "not_supported",
29 }
30 }
31}
32#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct ClaimResult {
35 pub verdict: Verdict,
37 pub probabilities: IndexMap<String, f64>,
39 pub confidence: f64,
41 pub automated: Automation,
43}
44#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
46#[serde(rename_all = "lowercase")]
47pub enum CompressionUnitKind {
48 #[default]
50 Sentence,
51 Line,
53 Paragraph,
55}
56#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct CompressionUnit {
59 pub start: usize,
61 pub end: usize,
63 pub probability: f64,
65}
66#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
68#[serde(rename_all = "lowercase")]
69pub enum ExtractSource {
70 Choice,
72 Span,
74}
75#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
77pub enum ExtractCalibration {
78 #[serde(rename = "validated")]
80 Validated,
81 #[serde(rename = "per decision")]
83 PerDecision,
84}
85#[derive(Debug, Clone, Serialize, Deserialize)]
87pub struct ExtractField {
88 pub value: Value,
90 pub confidence: f64,
92 pub source: ExtractSource,
94 pub calibration: ExtractCalibration,
96 pub probabilities: Option<IndexMap<String, f64>>,
98 pub present_probability: Option<f64>,
100 pub text: Option<String>,
102 pub start: Option<usize>,
104 pub end: Option<usize>,
106 pub line_probability: Option<f64>,
108 pub span_probability: Option<f64>,
110 pub error: Option<String>,
112}
113#[derive(Debug, Clone, Copy)]
115pub struct ClassifyAnswer<'a> {
116 pub scores: &'a IndexMap<String, f64>,
118 pub label: Option<&'a str>,
120 pub top_label: Option<&'a str>,
122 pub confidence: Option<f64>,
124 pub automated: Option<&'a Automation>,
126 pub labels: Option<&'a [String]>,
128 pub set_confidence: Option<f64>,
130}
131#[derive(Debug, Clone, Copy)]
133pub struct VerifyAnswer<'a> {
134 pub claims: &'a IndexMap<String, ClaimResult>,
136 pub all_supported: bool,
138 pub not_supported: &'a [String],
140}
141#[derive(Debug, Clone, Copy)]
143pub struct CompressAnswer<'a> {
144 pub compressed: &'a str,
146 pub kept: &'a [usize],
148 pub units: &'a [CompressionUnit],
150 pub unit: CompressionUnitKind,
152 pub original_chars: usize,
154 pub compressed_chars: usize,
156 pub ratio: f64,
158}
159#[derive(Debug, Clone, Copy)]
161pub struct ExtractAnswer<'a> {
162 pub data: &'a IndexMap<String, Value>,
164 pub fields: &'a IndexMap<String, ExtractField>,
166}
167fn required<T>(v: Option<T>) -> Result<T, Error> {
168 v.ok_or(Error::InvalidResponse("missing task answer field"))
169}
170fn probability(v: f64) -> Result<(), Error> {
171 if v.is_finite() && (0.0..=1.0).contains(&v) {
172 Ok(())
173 } else {
174 Err(Error::InvalidResponse("invalid task probability"))
175 }
176}
177fn probabilities(v: &IndexMap<String, f64>) -> Result<(), Error> {
178 v.values().try_for_each(|n| probability(*n))
179}
180macro_rules! task_kind {
181 ($marker:ident,$view:ident,$variant:ident,$a:ident,$body:block) => {
182 #[doc=concat!("Result marker for ",stringify!($variant)," tasks.")]
183 #[derive(Debug)]
184 pub struct $marker;
185 impl crate::questions::sealed::Sealed for $marker {}
186 impl AnswerKind for $marker {
187 type Output<'a> = $view<'a>;
188 fn parse($a: &Answer) -> Result<Self::Output<'_>, Error> {
189 if !matches!($a.kind, QuestionType::$variant) || $a.experimental != Some(true) {
190 return Err(Error::InvalidResponse("mismatched experimental task"));
191 }
192 $body
193 }
194 }
195 };
196}
197task_kind!(ClassifyKind, ClassifyAnswer, Classify, a, {
198 let scores = required(a.scores.as_ref())?;
199 probabilities(scores)?;
200 if let Some(labels) = &a.labels {
201 if labels.iter().any(|k| !scores.contains_key(k)) {
202 return Err(Error::InvalidResponse("unknown classification label"));
203 }
204 probability(required(a.set_confidence)?)?;
205 } else {
206 let label = required(a.label.as_deref())?;
207 if !scores.contains_key(label) || a.top_label.as_deref() != Some(label) {
208 return Err(Error::InvalidResponse("invalid classification label"));
209 }
210 probability(required(a.confidence)?)?;
211 required(a.automated.as_ref())?;
212 }
213 Ok(ClassifyAnswer {
214 scores,
215 label: a.label.as_deref(),
216 top_label: a.top_label.as_deref(),
217 confidence: a.confidence,
218 automated: a.automated.as_ref(),
219 labels: a.labels.as_deref(),
220 set_confidence: a.set_confidence,
221 })
222});
223task_kind!(VerifyKind, VerifyAnswer, Verify, a, {
224 let claims = required(a.claims.as_ref())?;
225 for c in claims.values() {
226 probabilities(&c.probabilities)?;
227 probability(c.confidence)?;
228 if !c.probabilities.contains_key(c.verdict.key()) {
229 return Err(Error::InvalidResponse("invalid claim verdict"));
230 }
231 }
232 let flagged = required(a.not_supported.as_deref())?;
233 if flagged.iter().any(|k| !claims.contains_key(k)) {
234 return Err(Error::InvalidResponse("invalid flagged claim"));
235 }
236 Ok(VerifyAnswer { claims, all_supported: required(a.all_supported)?, not_supported: flagged })
237});
238task_kind!(CompressKind, CompressAnswer, Compress, a, {
239 let units = required(a.units.as_deref())?;
240 for u in units {
241 probability(u.probability)?;
242 if u.end < u.start {
243 return Err(Error::InvalidResponse("reversed compression span"));
244 }
245 }
246 let kept = required(a.kept.as_deref())?;
247 if kept.iter().any(|i| *i >= units.len()) {
248 return Err(Error::InvalidResponse("invalid kept unit"));
249 }
250 let ratio = required(a.ratio)?;
251 if !ratio.is_finite() || ratio < 0.0 {
252 return Err(Error::InvalidResponse("invalid compression ratio"));
253 }
254 Ok(CompressAnswer {
255 compressed: required(a.compressed.as_deref())?,
256 kept,
257 units,
258 unit: required(a.unit)?,
259 original_chars: required(a.original_chars)?,
260 compressed_chars: required(a.compressed_chars)?,
261 ratio,
262 })
263});
264task_kind!(ExtractKind, ExtractAnswer, Extract, a, {
265 let fields = required(a.fields.as_ref())?;
266 for f in fields.values() {
267 probability(f.confidence)?;
268 if let Some(p) = &f.probabilities {
269 probabilities(p)?;
270 }
271 for p in [f.present_probability, f.line_probability, f.span_probability].into_iter().flatten() {
272 probability(p)?;
273 }
274 if f.text.is_some() && required(f.end)? < required(f.start)? {
275 return Err(Error::InvalidResponse("reversed extraction span"));
276 }
277 }
278 Ok(ExtractAnswer { data: required(a.data.as_ref())?, fields })
279});
280pub type ClassifyQuestion = QuestionSpec<ClassifyKind>;
282pub type VerifyQuestion = QuestionSpec<VerifyKind>;
284pub type CompressQuestion = QuestionSpec<CompressKind>;
286pub type ExtractQuestion = QuestionSpec<ExtractKind>;
288#[derive(Debug, Clone, Copy, Default)]
290pub struct ClassifyOptions<'a> {
291 pub multi_label: bool,
293 pub none: bool,
295 pub max_error: Option<&'a [f64]>,
297 pub hide_keys: bool,
299 pub selection: crate::MultiOptions<'a>,
301}
302#[derive(Debug, Clone, Copy, Default)]
304pub struct VerifyOptions<'a> {
305 pub binary: bool,
307 pub max_error: Option<&'a [f64]>,
309}
310#[derive(Debug, Clone, Copy, Default)]
312pub struct CompressOptions {
313 pub unit: CompressionUnitKind,
315 pub keep_ratio: Option<f64>,
317 pub threshold: Option<f64>,
319 pub min_keep: Option<usize>,
321}
322fn task(kind: QuestionType, instructions: &str) -> Question {
323 let mut q = Question::new(kind, instructions);
324 if instructions.is_empty() {
325 q.instructions = None;
326 }
327 q
328}
329fn risks(q: &mut Question, values: Option<&[f64]>) -> Result<(), Error> {
330 if let Some(values) = values {
331 if values.len() > 64 || values.iter().any(|v| !v.is_finite() || *v <= 0.0 || *v >= 1.0) {
332 return Err(Error::InvalidRequest("invalid task max_error"));
333 }
334 q.extra.insert("max_error".into(), serde_json::json!(values));
335 }
336 Ok(())
337}
338pub fn classify(instructions: &str, labels: Value, options: ClassifyOptions<'_>) -> Result<ClassifyQuestion, Error> {
340 let count = match &labels {
341 Value::Object(m) => m.len(),
342 Value::Array(a) => a.len(),
343 _ => 0,
344 };
345 if !(1..=256).contains(&count) {
346 return Err(Error::InvalidRequest("classify requires 1..256 labels"));
347 }
348 let mut q = task(QuestionType::Classify, instructions);
349 q.extra.insert("labels".into(), labels);
350 q.extra.insert("multi_label".into(), options.multi_label.into());
351 q.extra.insert("none".into(), options.none.into());
352 risks(&mut q, options.max_error)?;
353 if options.hide_keys {
354 q.extra.insert("option_keys".into(), "hide".into());
355 }
356 if options.multi_label {
357 if !options.selection.exclusive.is_empty()
358 || options.selection.fallback.is_some()
359 || options.selection.fallback_threshold.is_some()
360 {
361 return Err(Error::InvalidRequest("classify does not support exclusive/fallback"));
362 }
363 q.extra.extend(crate::questions::selection_fields(options.selection, None)?);
364 } else if options.selection.threshold.is_some()
365 || options.selection.min.is_some()
366 || options.selection.max.is_some()
367 || !matches!(options.selection.objective, MultiObjective::Threshold)
368 || !options.selection.exclusive.is_empty()
369 || options.selection.fallback.is_some()
370 || options.selection.fallback_threshold.is_some()
371 {
372 return Err(Error::InvalidRequest("selection options require multi_label"));
373 }
374 Ok(QuestionSpec::new(q))
375}
376pub fn verify(instructions: &str, claims: Value, options: VerifyOptions<'_>) -> Result<VerifyQuestion, Error> {
378 if match &claims {
379 Value::Object(m) => m.is_empty(),
380 Value::Array(a) => a.is_empty(),
381 _ => true,
382 } {
383 return Err(Error::InvalidRequest("verify requires claims"));
384 }
385 let mut q = task(QuestionType::Verify, instructions);
386 q.extra.insert("claims".into(), claims);
387 q.extra.insert("binary".into(), options.binary.into());
388 risks(&mut q, options.max_error)?;
389 Ok(QuestionSpec::new(q))
390}
391pub fn compress(instructions: &str, options: CompressOptions) -> Result<CompressQuestion, Error> {
393 let mut q = task(QuestionType::Compress, instructions);
394 q.extra.insert("unit".into(), serde_json::to_value(options.unit).map_err(Error::Encode)?);
395 for (k, v) in [("keep_ratio", options.keep_ratio), ("threshold", options.threshold)] {
396 if let Some(v) = v {
397 if !v.is_finite() || !(0.0..=1.0).contains(&v) || (k == "keep_ratio" && v == 0.0) {
398 return Err(Error::InvalidRequest("invalid compression probability"));
399 }
400 q.extra.insert(k.into(), v.into());
401 }
402 }
403 if let Some(n) = options.min_keep {
404 q.extra.insert("min_keep".into(), n.into());
405 }
406 Ok(QuestionSpec::new(q))
407}
408pub fn extract(instructions: &str, schema: Value) -> Result<ExtractQuestion, Error> {
410 if !schema.is_object() {
411 return Err(Error::InvalidRequest("schema must be an object"));
412 }
413 let mut q = task(QuestionType::Extract, instructions);
414 q.extra.insert("schema".into(), schema);
415 Ok(QuestionSpec::new(q))
416}