Skip to main content

basal_client/
tasks.rs

1//! Experimental Basal 2.1 tasks. Source offsets count Unicode scalar values, not UTF-8 bytes.
2use crate::{Answer, AnswerKind, Error, MultiObjective, Question, QuestionSpec, QuestionType};
3use indexmap::IndexMap;
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7/// Certified automation by requested error target. `None` is an unavailable task certificate.
8pub type Automation = IndexMap<String, Option<bool>>;
9/// Claim verification verdict.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum Verdict {
13    /// Supported by the source.
14    Supported,
15    /// Contradicted by the source.
16    Contradicted,
17    /// Not stated in the source.
18    NotStated,
19    /// Binary verification did not support the claim.
20    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/// Per-claim calibrated verdict.
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct ClaimResult {
35    /// Verdict.
36    pub verdict: Verdict,
37    /// Outcome probabilities.
38    pub probabilities: IndexMap<String, f64>,
39    /// Maximum outcome probability.
40    pub confidence: f64,
41    /// Certification for this task.
42    pub automated: Automation,
43}
44/// Compression segmentation.
45#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
46#[serde(rename_all = "lowercase")]
47pub enum CompressionUnitKind {
48    /// Sentence segmentation.
49    #[default]
50    Sentence,
51    /// Nonempty lines.
52    Line,
53    /// Paragraphs separated by blank lines.
54    Paragraph,
55}
56/// One compression unit in the original text.
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct CompressionUnit {
59    /// Inclusive Unicode code point offset.
60    pub start: usize,
61    /// Exclusive Unicode code point offset.
62    pub end: usize,
63    /// Retention probability.
64    pub probability: f64,
65}
66/// Extraction decision source.
67#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
68#[serde(rename_all = "lowercase")]
69pub enum ExtractSource {
70    /// Enum/boolean decision.
71    Choice,
72    /// Source span selection.
73    Span,
74}
75/// Extraction calibration scope; span products are not certified field probabilities.
76#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
77pub enum ExtractCalibration {
78    /// Categorical decision.
79    #[serde(rename = "validated")]
80    Validated,
81    /// Each constituent decision is calibrated separately.
82    #[serde(rename = "per decision")]
83    PerDecision,
84}
85/// Extracted field, including missing values and source provenance.
86#[derive(Debug, Clone, Serialize, Deserialize)]
87pub struct ExtractField {
88    /// Native JSON value, null when missing or numeric parsing failed.
89    pub value: Value,
90    /// Decision confidence or line/span product.
91    pub confidence: f64,
92    /// Choice or span selection.
93    pub source: ExtractSource,
94    /// Calibration scope.
95    pub calibration: ExtractCalibration,
96    /// Categorical probabilities.
97    pub probabilities: Option<IndexMap<String, f64>>,
98    /// Probability the field is stated.
99    pub present_probability: Option<f64>,
100    /// Verbatim source text, when selected.
101    pub text: Option<String>,
102    /// Inclusive Unicode code point offset.
103    pub start: Option<usize>,
104    /// Exclusive Unicode code point offset.
105    pub end: Option<usize>,
106    /// Line selection probability.
107    pub line_probability: Option<f64>,
108    /// Span selection probability.
109    pub span_probability: Option<f64>,
110    /// Numeric conversion diagnostic.
111    pub error: Option<String>,
112}
113/// Checked single- or multi-label classification.
114#[derive(Debug, Clone, Copy)]
115pub struct ClassifyAnswer<'a> {
116    /// Label scores.
117    pub scores: &'a IndexMap<String, f64>,
118    /// Single selected label; absent in multi mode.
119    pub label: Option<&'a str>,
120    /// Top label; absent in multi mode.
121    pub top_label: Option<&'a str>,
122    /// Single-label confidence.
123    pub confidence: Option<f64>,
124    /// Task certification; absent in multi mode.
125    pub automated: Option<&'a Automation>,
126    /// Multi-label set; absent in single mode.
127    pub labels: Option<&'a [String]>,
128    /// Multi-label set confidence.
129    pub set_confidence: Option<f64>,
130}
131/// Checked claim verification.
132#[derive(Debug, Clone, Copy)]
133pub struct VerifyAnswer<'a> {
134    /// Per-claim results.
135    pub claims: &'a IndexMap<String, ClaimResult>,
136    /// Whether every claim is supported.
137    pub all_supported: bool,
138    /// Flagged claim keys.
139    pub not_supported: &'a [String],
140}
141/// Checked extractive compression.
142#[derive(Debug, Clone, Copy)]
143pub struct CompressAnswer<'a> {
144    /// Compressed source text.
145    pub compressed: &'a str,
146    /// Kept source unit indices.
147    pub kept: &'a [usize],
148    /// Source spans and retention probabilities.
149    pub units: &'a [CompressionUnit],
150    /// Segmentation unit.
151    pub unit: CompressionUnitKind,
152    /// Original Unicode code point count.
153    pub original_chars: usize,
154    /// Compressed Unicode code point count.
155    pub compressed_chars: usize,
156    /// Character ratio.
157    pub ratio: f64,
158}
159/// Checked structured extraction.
160#[derive(Debug, Clone, Copy)]
161pub struct ExtractAnswer<'a> {
162    /// Nested output data.
163    pub data: &'a IndexMap<String, Value>,
164    /// Field provenance and calibrated constituent decisions.
165    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});
280/// Classification question definition.
281pub type ClassifyQuestion = QuestionSpec<ClassifyKind>;
282/// Verification question definition.
283pub type VerifyQuestion = QuestionSpec<VerifyKind>;
284/// Compression question definition.
285pub type CompressQuestion = QuestionSpec<CompressKind>;
286/// Extraction question definition.
287pub type ExtractQuestion = QuestionSpec<ExtractKind>;
288/// Experimental classification options.
289#[derive(Debug, Clone, Copy, Default)]
290pub struct ClassifyOptions<'a> {
291    /// Independent multi-label classification.
292    pub multi_label: bool,
293    /// Add a none-of-the-above label in single mode.
294    pub none: bool,
295    /// Requested certified error targets; omitted preserves server defaults.
296    pub max_error: Option<&'a [f64]>,
297    /// Hide category names in the decision prompt.
298    pub hide_keys: bool,
299    /// Multi-label selection options.
300    pub selection: crate::MultiOptions<'a>,
301}
302/// Experimental verification options.
303#[derive(Debug, Clone, Copy, Default)]
304pub struct VerifyOptions<'a> {
305    /// Supported/not-supported instead of three outcomes.
306    pub binary: bool,
307    /// Requested certified error targets.
308    pub max_error: Option<&'a [f64]>,
309}
310/// Extractive compression options.
311#[derive(Debug, Clone, Copy, Default)]
312pub struct CompressOptions {
313    /// Segmentation, sentences by default.
314    pub unit: CompressionUnitKind,
315    /// Character retention target, whole units may exceed it.
316    pub keep_ratio: Option<f64>,
317    /// Independent retention threshold.
318    pub threshold: Option<f64>,
319    /// Minimum retained units; omitted uses server default 1.
320    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}
338/// Classify using a label object or ordered list, normalized by the server.
339pub 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}
376/// Verify an ordered map or list of claims.
377pub 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}
391/// Compress verbatim source units; the model chooses units, it does not generate text.
392pub 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}
408/// Extract enum, boolean, scalar and nested object fields from a JSON schema.
409pub 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}