laminar_db/ai/
adapter.rs

1//! Adapt a raw provider response to a task's per-row output. Stays Arrow-free;
2//! the operator builds the column from the returned [`InferenceOutputs`].
3
4use thiserror::Error;
5
6use crate::ai::provider::InferenceOutputs;
7use crate::ai::registry::{BackendKind, Task};
8
9/// Errors from adapting a response to a task output.
10#[derive(Debug, Error, PartialEq, Eq)]
11pub enum AdapterError {
12    /// A local classifier produced logits but no labels were supplied to map
13    /// them to.
14    #[error("classification requires the model's labels but none were provided")]
15    MissingLabels,
16
17    /// argmax landed on an index with no corresponding label.
18    #[error("classifier chose index {index} but only {len} labels are defined")]
19    LabelIndexOutOfRange {
20        /// The chosen index.
21        index: usize,
22        /// Number of available labels.
23        len: usize,
24    },
25
26    /// A classifier returned an empty logit vector for a row.
27    #[error("classifier returned no logits for a row")]
28    EmptyLogits,
29
30    /// Sentiment scoring needs both a `positive` and a `negative` label.
31    #[error(
32        "sentiment scoring requires 'positive' and 'negative' among the labels, got {labels:?}"
33    )]
34    SentimentLabelsUnusable {
35        /// Labels that were available.
36        labels: Vec<String>,
37    },
38
39    /// A remote sentiment reply contained no parseable number.
40    #[error("remote sentiment reply had no parseable score: {reply:?}")]
41    UnparseableScore {
42        /// The unparseable reply string.
43        reply: String,
44    },
45
46    /// The raw output shape did not match what the task/backend produces.
47    #[error("task '{task}' on the {kind:?} backend expected {expected} output, got {got}")]
48    UnexpectedOutputShape {
49        /// The task.
50        task: Task,
51        /// The backend kind.
52        kind: BackendKind,
53        /// Expected shape name.
54        expected: &'static str,
55        /// Actual shape name.
56        got: &'static str,
57    },
58
59    /// The task cannot run on this backend kind (e.g. generation on local).
60    #[error("task '{task}' is not supported on the {kind:?} backend")]
61    UnsupportedCombination {
62        /// The task.
63        task: Task,
64        /// The backend kind.
65        kind: BackendKind,
66    },
67}
68
69/// Map a raw provider response to the task's per-row output.
70///
71/// `labels` is the candidate/intrinsic label set for classification; required
72/// for local classifiers, ignored otherwise.
73///
74/// # Errors
75///
76/// Returns [`AdapterError`] if the output shape is wrong, a local classifier
77/// has no labels or an out-of-range argmax, or the task/backend combination is
78/// unsupported.
79pub fn parse_response(
80    task: Task,
81    kind: BackendKind,
82    raw: InferenceOutputs,
83    labels: Option<&[String]>,
84) -> Result<InferenceOutputs, AdapterError> {
85    match task {
86        Task::Classify => match kind {
87            BackendKind::Local => classify_from_logits(kind, raw, labels),
88            BackendKind::Remote => coerce_remote_classification(raw, labels),
89        },
90        Task::Sentiment => match kind {
91            BackendKind::Local => sentiment_from_logits(raw, labels),
92            BackendKind::Remote => sentiment_from_text(raw),
93        },
94        Task::Embed => expect_vectors(task, kind, raw),
95        Task::Complete | Task::Summarize | Task::Translate | Task::Gen | Task::Extract => {
96            if kind != BackendKind::Remote {
97                return Err(AdapterError::UnsupportedCombination { task, kind });
98            }
99            expect_text(task, kind, raw)
100        }
101    }
102}
103
104/// argmax over each row's logits, mapped to labels.
105fn classify_from_logits(
106    kind: BackendKind,
107    raw: InferenceOutputs,
108    labels: Option<&[String]>,
109) -> Result<InferenceOutputs, AdapterError> {
110    let rows = match raw {
111        InferenceOutputs::Vectors(rows) => rows,
112        other => {
113            return Err(AdapterError::UnexpectedOutputShape {
114                task: Task::Classify,
115                kind,
116                expected: "vectors",
117                got: shape_name(&other),
118            });
119        }
120    };
121    let labels = labels.ok_or(AdapterError::MissingLabels)?;
122    let mut out = Vec::with_capacity(rows.len());
123    for logits in rows {
124        let index = argmax(&logits).ok_or(AdapterError::EmptyLogits)?;
125        let label = labels
126            .get(index)
127            .ok_or(AdapterError::LabelIndexOutOfRange {
128                index,
129                len: labels.len(),
130            })?;
131        out.push(label.clone());
132    }
133    Ok(InferenceOutputs::Text(out))
134}
135
136/// Normalize each reply to a canonical label (exact then containment match),
137/// or fall back to trimmed raw text if none matches.
138fn coerce_remote_classification(
139    raw: InferenceOutputs,
140    labels: Option<&[String]>,
141) -> Result<InferenceOutputs, AdapterError> {
142    let texts = match raw {
143        InferenceOutputs::Text(texts) => texts,
144        other => {
145            return Err(AdapterError::UnexpectedOutputShape {
146                task: Task::Classify,
147                kind: BackendKind::Remote,
148                expected: "text",
149                got: shape_name(&other),
150            });
151        }
152    };
153    let coerced = texts
154        .into_iter()
155        .map(|text| match labels {
156            Some(labels) => coerce_label(&text, labels),
157            None => text.trim().to_string(),
158        })
159        .collect();
160    Ok(InferenceOutputs::Text(coerced))
161}
162
163/// Exact (case-insensitive) then containment match; falls back to trimmed text.
164fn coerce_label(text: &str, labels: &[String]) -> String {
165    let trimmed = text.trim();
166    if let Some(label) = labels.iter().find(|l| l.eq_ignore_ascii_case(trimmed)) {
167        return label.clone();
168    }
169    let lower = trimmed.to_ascii_lowercase();
170    if let Some(label) = labels
171        .iter()
172        .find(|l| lower.contains(&l.to_ascii_lowercase()))
173    {
174        return label.clone();
175    }
176    trimmed.to_string()
177}
178
179/// Pass text outputs through unchanged, rejecting other shapes.
180fn expect_text(
181    task: Task,
182    kind: BackendKind,
183    raw: InferenceOutputs,
184) -> Result<InferenceOutputs, AdapterError> {
185    match raw {
186        InferenceOutputs::Text(text) => Ok(InferenceOutputs::Text(text)),
187        other => Err(AdapterError::UnexpectedOutputShape {
188            task,
189            kind,
190            expected: "text",
191            got: shape_name(&other),
192        }),
193    }
194}
195
196/// Pass vector outputs through unchanged, rejecting other shapes.
197fn expect_vectors(
198    task: Task,
199    kind: BackendKind,
200    raw: InferenceOutputs,
201) -> Result<InferenceOutputs, AdapterError> {
202    match raw {
203        InferenceOutputs::Vectors(vectors) => Ok(InferenceOutputs::Vectors(vectors)),
204        other => Err(AdapterError::UnexpectedOutputShape {
205            task,
206            kind,
207            expected: "vectors",
208            got: shape_name(&other),
209        }),
210    }
211}
212
213/// Static shape name for error messages.
214fn shape_name(raw: &InferenceOutputs) -> &'static str {
215    match raw {
216        InferenceOutputs::Text(_) => "text",
217        InferenceOutputs::Vectors(_) => "vectors",
218        InferenceOutputs::Scores(_) => "scores",
219    }
220}
221
222/// Softmax each row's logits, return `P(positive) − P(negative)` in `[-1, 1]`.
223fn sentiment_from_logits(
224    raw: InferenceOutputs,
225    labels: Option<&[String]>,
226) -> Result<InferenceOutputs, AdapterError> {
227    let rows = match raw {
228        InferenceOutputs::Vectors(rows) => rows,
229        other => {
230            return Err(AdapterError::UnexpectedOutputShape {
231                task: Task::Sentiment,
232                kind: BackendKind::Local,
233                expected: "vectors",
234                got: shape_name(&other),
235            });
236        }
237    };
238    let labels = labels.ok_or(AdapterError::MissingLabels)?;
239    let pos = labels
240        .iter()
241        .position(|l| l.eq_ignore_ascii_case("positive"));
242    let neg = labels
243        .iter()
244        .position(|l| l.eq_ignore_ascii_case("negative"));
245    let (Some(pos), Some(neg)) = (pos, neg) else {
246        return Err(AdapterError::SentimentLabelsUnusable {
247            labels: labels.to_vec(),
248        });
249    };
250    let mut scores = Vec::with_capacity(rows.len());
251    for logits in rows {
252        if logits.is_empty() {
253            return Err(AdapterError::EmptyLogits);
254        }
255        let probs = softmax(&logits);
256        let p_pos = probs.get(pos).copied().unwrap_or(0.0);
257        let p_neg = probs.get(neg).copied().unwrap_or(0.0);
258        scores.push(f64::from(p_pos - p_neg));
259    }
260    Ok(InferenceOutputs::Scores(scores))
261}
262
263/// Parse a number from each reply and clamp to `[-1, 1]`.
264fn sentiment_from_text(raw: InferenceOutputs) -> Result<InferenceOutputs, AdapterError> {
265    let texts = match raw {
266        InferenceOutputs::Text(texts) => texts,
267        other => {
268            return Err(AdapterError::UnexpectedOutputShape {
269                task: Task::Sentiment,
270                kind: BackendKind::Remote,
271                expected: "text",
272                got: shape_name(&other),
273            });
274        }
275    };
276    let mut scores = Vec::with_capacity(texts.len());
277    for reply in texts {
278        let value = parse_score(&reply).ok_or(AdapterError::UnparseableScore { reply })?;
279        scores.push(value.clamp(-1.0, 1.0));
280    }
281    Ok(InferenceOutputs::Scores(scores))
282}
283
284/// Numerically stable softmax.
285fn softmax(logits: &[f32]) -> Vec<f32> {
286    let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
287    let mut exps: Vec<f32> = logits.iter().map(|&v| (v - max).exp()).collect();
288    let sum: f32 = exps.iter().sum();
289    if sum > 0.0 {
290        for e in &mut exps {
291            *e /= sum;
292        }
293    }
294    exps
295}
296
297/// Extract the first number from a reply string, or `None`.
298fn parse_score(reply: &str) -> Option<f64> {
299    let trimmed = reply.trim();
300    if let Ok(v) = trimmed.parse::<f64>() {
301        return Some(v);
302    }
303    trimmed.split_whitespace().find_map(|token| {
304        // Strip punctuation: leading sign/digit may start; only a digit may end
305        // (so a sentence-final '.' isn't kept).
306        token
307            .trim_start_matches(|c: char| !(c.is_ascii_digit() || c == '-'))
308            .trim_end_matches(|c: char| !c.is_ascii_digit())
309            .parse::<f64>()
310            .ok()
311    })
312}
313
314/// Index of the maximum value, or `None` if empty.
315fn argmax(values: &[f32]) -> Option<usize> {
316    let mut best: Option<(usize, f32)> = None;
317    for (index, &value) in values.iter().enumerate() {
318        match best {
319            Some((_, current)) if value <= current => {}
320            _ => best = Some((index, value)),
321        }
322    }
323    best.map(|(index, _)| index)
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329
330    fn labels() -> Vec<String> {
331        vec!["negative".into(), "positive".into(), "neutral".into()]
332    }
333
334    #[test]
335    fn local_classify_argmaxes_logits_to_labels() {
336        let raw = InferenceOutputs::Vectors(vec![vec![0.1, 0.9, 0.2], vec![5.0, 0.0, 0.0]]);
337        let out = parse_response(Task::Classify, BackendKind::Local, raw, Some(&labels())).unwrap();
338        assert_eq!(
339            out,
340            InferenceOutputs::Text(vec!["positive".into(), "negative".into()])
341        );
342    }
343
344    #[test]
345    fn local_classify_requires_labels() {
346        let raw = InferenceOutputs::Vectors(vec![vec![0.1, 0.9]]);
347        assert_eq!(
348            parse_response(Task::Sentiment, BackendKind::Local, raw, None),
349            Err(AdapterError::MissingLabels)
350        );
351    }
352
353    #[test]
354    fn remote_classify_coerces_to_canonical_label() {
355        // Exact (case-insensitive), containment in a wrapped reply, and the
356        // no-match fallback to trimmed raw text.
357        let raw = InferenceOutputs::Text(vec![
358            "The sentiment is Positive.".into(),
359            "NEGATIVE".into(),
360            "totally unrelated".into(),
361        ]);
362        let out =
363            parse_response(Task::Classify, BackendKind::Remote, raw, Some(&labels())).unwrap();
364        assert_eq!(
365            out,
366            InferenceOutputs::Text(vec![
367                "positive".into(),
368                "negative".into(),
369                "totally unrelated".into(),
370            ])
371        );
372    }
373
374    #[test]
375    fn local_sentiment_softmaxes_logits_to_a_signed_score() {
376        // labels = [negative, positive, neutral]. Row 1 favours positive, row 2
377        // favours negative. Score = P(pos) − P(neg) ∈ [-1, 1].
378        let raw = InferenceOutputs::Vectors(vec![vec![0.0, 2.0, 0.0], vec![3.0, 0.0, 0.0]]);
379        let out =
380            parse_response(Task::Sentiment, BackendKind::Local, raw, Some(&labels())).unwrap();
381        let InferenceOutputs::Scores(scores) = out else {
382            panic!("sentiment is numeric");
383        };
384        // Row 1: softmax([0,2,0]) → P(neg)=P(neu)=e^0/(2+e^2), P(pos)=e^2/(2+e^2).
385        let denom = 2.0 + std::f32::consts::E.powi(2);
386        let expected0 = f64::from((std::f32::consts::E.powi(2) - 1.0) / denom);
387        assert!((scores[0] - expected0).abs() < 1e-6, "got {}", scores[0]);
388        assert!(scores[0] > 0.0 && scores[1] < 0.0);
389        assert!((-1.0..=1.0).contains(&scores[0]) && (-1.0..=1.0).contains(&scores[1]));
390    }
391
392    #[test]
393    fn local_sentiment_needs_positive_and_negative_labels() {
394        let raw = InferenceOutputs::Vectors(vec![vec![0.1, 0.9]]);
395        let stars = vec!["one_star".to_string(), "five_star".to_string()];
396        assert!(matches!(
397            parse_response(Task::Sentiment, BackendKind::Local, raw, Some(&stars)),
398            Err(AdapterError::SentimentLabelsUnusable { .. })
399        ));
400    }
401
402    #[test]
403    fn remote_sentiment_parses_and_clamps_a_number() {
404        let raw = InferenceOutputs::Text(vec![
405            "0.8".into(),
406            "The sentiment is -0.5.".into(),
407            "2.0".into(), // out of range → clamped
408        ]);
409        let out = parse_response(Task::Sentiment, BackendKind::Remote, raw, None).unwrap();
410        let InferenceOutputs::Scores(scores) = out else {
411            panic!("sentiment is numeric");
412        };
413        assert!((scores[0] - 0.8).abs() < 1e-12);
414        assert!((scores[1] + 0.5).abs() < 1e-12);
415        assert!((scores[2] - 1.0).abs() < 1e-12);
416    }
417
418    #[test]
419    fn generation_is_remote_only() {
420        let raw = InferenceOutputs::Vectors(vec![vec![0.0]]);
421        assert_eq!(
422            parse_response(Task::Complete, BackendKind::Local, raw, None),
423            Err(AdapterError::UnsupportedCombination {
424                task: Task::Complete,
425                kind: BackendKind::Local,
426            })
427        );
428    }
429}
laminar_db/ai/adapter.rs

laminar_db/ai/
adapter.rs
