sie_sdk/types/
item.rs

1//! Request items: the multimodal unit every inference endpoint consumes.
2
3use std::path::{Path, PathBuf};
4
5use rmpv::Value as MsgValue;
6use serde_json::Value;
7
8use crate::error::{Error, Result};
9use crate::media::{self, Samples};
10
11/// An image, in whatever form the caller already has it.
12#[derive(Debug, Clone, PartialEq)]
13enum ImageData {
14    /// Already encoded; the format is sniffed from the bytes.
15    Encoded(Vec<u8>),
16    /// Read from disk at send time, then sniffed.
17    Path(PathBuf),
18    /// Decoded pixels, re-encoded as JPEG at send time.
19    Decoded(Box<image::DynamicImage>),
20}
21
22/// One image attached to an [`Item`].
23#[derive(Debug, Clone, PartialEq)]
24pub struct ImageInput {
25    data: ImageData,
26    format: Option<String>,
27}
28
29impl ImageInput {
30    /// An image that is already encoded. The format is detected from its magic bytes.
31    pub fn bytes(data: impl Into<Vec<u8>>) -> Self {
32        Self {
33            data: ImageData::Encoded(data.into()),
34            format: None,
35        }
36    }
37
38    /// An encoded image read from disk when the request is sent.
39    pub fn path(path: impl Into<PathBuf>) -> Self {
40        Self {
41            data: ImageData::Path(path.into()),
42            format: None,
43        }
44    }
45
46    /// Decoded pixels, re-encoded as JPEG when the request is sent.
47    pub fn decoded(image: image::DynamicImage) -> Self {
48        Self {
49            data: ImageData::Decoded(Box::new(image)),
50            format: None,
51        }
52    }
53
54    /// Declare the format explicitly. A declaration that contradicts the bytes is an error.
55    pub fn format(mut self, format: impl Into<String>) -> Self {
56        self.format = Some(format.into());
57        self
58    }
59
60    pub(crate) fn resolve(&self) -> Result<(Vec<u8>, String)> {
61        let encoded = match &self.data {
62            ImageData::Encoded(data) => data.clone(),
63            ImageData::Path(path) => read_file(path)?,
64            ImageData::Decoded(image) => media::encode_jpeg(image)?,
65        };
66        let detected = media::detect_image_format(&encoded)?;
67        if let Some(declared) = &self.format {
68            let declared = media::canonical_format(declared)?;
69            if declared != detected {
70                return Err(Error::invalid(format!(
71                    "Image format mismatch: declared {declared:?}, detected {detected:?}"
72                )));
73            }
74        }
75        Ok((encoded, detected.to_string()))
76    }
77}
78
79#[derive(Debug, Clone, PartialEq)]
80enum AudioData {
81    Encoded(Vec<u8>),
82    Path(PathBuf),
83    Waveform {
84        samples: Samples,
85        channels: u16,
86        sample_rate: u32,
87    },
88}
89
90/// Audio attached to an [`Item`].
91#[derive(Debug, Clone, PartialEq)]
92pub struct AudioInput {
93    data: AudioData,
94    format: Option<String>,
95}
96
97impl AudioInput {
98    /// Audio that is already in a container the server understands.
99    pub fn bytes(data: impl Into<Vec<u8>>) -> Self {
100        Self {
101            data: AudioData::Encoded(data.into()),
102            format: None,
103        }
104    }
105
106    /// Encoded audio read from disk when the request is sent; the format comes from the
107    /// filename.
108    pub fn path(path: impl Into<PathBuf>) -> Self {
109        Self {
110            data: AudioData::Path(path.into()),
111            format: None,
112        }
113    }
114
115    /// A raw waveform, wrapped in a 16-bit PCM WAV container when the request is sent.
116    pub fn waveform(samples: Samples, channels: u16, sample_rate: u32) -> Self {
117        Self {
118            data: AudioData::Waveform {
119                samples,
120                channels,
121                sample_rate,
122            },
123            format: None,
124        }
125    }
126
127    /// Declare the container format explicitly.
128    pub fn format(mut self, format: impl Into<String>) -> Self {
129        self.format = Some(format.into());
130        self
131    }
132
133    pub(crate) fn resolve(&self) -> Result<(Vec<u8>, Option<String>, Option<u32>)> {
134        match &self.data {
135            AudioData::Encoded(data) => Ok((data.clone(), self.format.clone(), None)),
136            AudioData::Path(path) => {
137                let inferred = media::infer_audio_format(path).map(str::to_string);
138                Ok((read_file(path)?, self.format.clone().or(inferred), None))
139            }
140            AudioData::Waveform {
141                samples,
142                channels,
143                sample_rate,
144            } => Ok((
145                media::encode_wav(samples, *channels, *sample_rate)?,
146                Some(self.format.clone().unwrap_or_else(|| "wav".to_string())),
147                Some(*sample_rate),
148            )),
149        }
150    }
151}
152
153/// A document or video attached to an [`Item`]: bytes plus an optional format token.
154#[derive(Debug, Clone, PartialEq)]
155pub struct BinaryInput {
156    data: BinaryData,
157    format: Option<String>,
158    kind: BinaryKind,
159}
160
161#[derive(Debug, Clone, PartialEq)]
162enum BinaryData {
163    Encoded(Vec<u8>),
164    Path(PathBuf),
165}
166
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
168enum BinaryKind {
169    Document,
170    Video,
171}
172
173impl BinaryInput {
174    /// A document held in memory.
175    pub fn document_bytes(data: impl Into<Vec<u8>>) -> Self {
176        Self {
177            data: BinaryData::Encoded(data.into()),
178            format: None,
179            kind: BinaryKind::Document,
180        }
181    }
182
183    /// A document read from disk; the format comes from the filename.
184    pub fn document_path(path: impl Into<PathBuf>) -> Self {
185        Self {
186            data: BinaryData::Path(path.into()),
187            format: None,
188            kind: BinaryKind::Document,
189        }
190    }
191
192    /// A video held in memory.
193    pub fn video_bytes(data: impl Into<Vec<u8>>) -> Self {
194        Self {
195            data: BinaryData::Encoded(data.into()),
196            format: None,
197            kind: BinaryKind::Video,
198        }
199    }
200
201    /// A video read from disk; the format comes from the filename.
202    pub fn video_path(path: impl Into<PathBuf>) -> Self {
203        Self {
204            data: BinaryData::Path(path.into()),
205            format: None,
206            kind: BinaryKind::Video,
207        }
208    }
209
210    /// Declare the format explicitly rather than inferring it from the filename.
211    pub fn format(mut self, format: impl Into<String>) -> Self {
212        self.format = Some(format.into());
213        self
214    }
215
216    pub(crate) fn resolve(&self) -> Result<(Vec<u8>, Option<String>)> {
217        match &self.data {
218            BinaryData::Encoded(data) => Ok((data.clone(), self.format.clone())),
219            BinaryData::Path(path) => {
220                let inferred = match self.kind {
221                    BinaryKind::Document => media::infer_document_format(path),
222                    BinaryKind::Video => media::infer_video_format(path),
223                }
224                .map(str::to_string);
225                Ok((read_file(path)?, self.format.clone().or(inferred)))
226            }
227        }
228    }
229}
230
231fn read_file(path: &Path) -> Result<Vec<u8>> {
232    std::fs::read(path).map_err(|err| {
233        Error::Io(std::io::Error::new(
234            err.kind(),
235            format!("could not read {}: {err}", path.display()),
236        ))
237    })
238}
239
240/// One unit of input: text, images, audio, video or a document, in any combination the
241/// model accepts.
242#[derive(Debug, Clone, Default, PartialEq)]
243#[allow(missing_docs)]
244pub struct Item {
245    /// Caller-chosen identifier, echoed back on the matching result.
246    pub id: Option<String>,
247    pub text: Option<String>,
248    pub images: Vec<ImageInput>,
249    pub audio: Option<AudioInput>,
250    pub video: Option<BinaryInput>,
251    pub document: Option<BinaryInput>,
252    /// Opaque metadata passed through to the model adapter.
253    pub metadata: Option<Value>,
254}
255
256impl Item {
257    /// An empty item.
258    pub fn new() -> Self {
259        Self::default()
260    }
261
262    /// A text-only item.
263    pub fn text(text: impl Into<String>) -> Self {
264        Self {
265            text: Some(text.into()),
266            ..Self::default()
267        }
268    }
269
270    /// An item holding a single image.
271    pub fn image(image: ImageInput) -> Self {
272        Self {
273            images: vec![image],
274            ..Self::default()
275        }
276    }
277
278    /// Set the identifier echoed back on the result.
279    pub fn with_id(mut self, id: impl Into<String>) -> Self {
280        self.id = Some(id.into());
281        self
282    }
283
284    /// Set or replace the text.
285    pub fn with_text(mut self, text: impl Into<String>) -> Self {
286        self.text = Some(text.into());
287        self
288    }
289
290    /// Append an image.
291    pub fn with_image(mut self, image: ImageInput) -> Self {
292        self.images.push(image);
293        self
294    }
295
296    /// Attach audio.
297    pub fn with_audio(mut self, audio: AudioInput) -> Self {
298        self.audio = Some(audio);
299        self
300    }
301
302    /// Attach a video.
303    pub fn with_video(mut self, video: BinaryInput) -> Self {
304        self.video = Some(video);
305        self
306    }
307
308    /// Attach a document.
309    pub fn with_document(mut self, document: BinaryInput) -> Self {
310        self.document = Some(document);
311        self
312    }
313
314    /// Attach adapter metadata.
315    pub fn with_metadata(mut self, metadata: Value) -> Self {
316        self.metadata = Some(metadata);
317        self
318    }
319
320    /// Encode to the msgpack shape, resolving every attachment to bytes.
321    ///
322    /// Only fields the caller set are emitted: `/v1` is additive, so an absent key and a
323    /// null one are not the same thing.
324    pub(crate) fn to_msgpack(&self) -> Result<MsgValue> {
325        let mut fields: Vec<(MsgValue, MsgValue)> = Vec::new();
326
327        if let Some(id) = &self.id {
328            fields.push((MsgValue::from("id"), MsgValue::from(id.as_str())));
329        }
330        if let Some(text) = &self.text {
331            fields.push((MsgValue::from("text"), MsgValue::from(text.as_str())));
332        }
333        if !self.images.is_empty() {
334            let mut images = Vec::with_capacity(self.images.len());
335            for image in &self.images {
336                let (data, format) = image.resolve()?;
337                images.push(MsgValue::Map(vec![
338                    (MsgValue::from("data"), MsgValue::Binary(data)),
339                    (MsgValue::from("format"), MsgValue::from(format)),
340                ]));
341            }
342            fields.push((MsgValue::from("images"), MsgValue::Array(images)));
343        }
344        if let Some(audio) = &self.audio {
345            let (data, format, sample_rate) = audio.resolve()?;
346            fields.push((
347                MsgValue::from("audio"),
348                MsgValue::Map(vec![
349                    (MsgValue::from("data"), MsgValue::Binary(data)),
350                    (MsgValue::from("format"), optional_str(format)),
351                    (
352                        MsgValue::from("sample_rate"),
353                        sample_rate.map_or(MsgValue::Nil, |rate| MsgValue::from(u64::from(rate))),
354                    ),
355                ]),
356            ));
357        }
358        if let Some(video) = &self.video {
359            let (data, format) = video.resolve()?;
360            fields.push((
361                MsgValue::from("video"),
362                MsgValue::Map(vec![
363                    (MsgValue::from("data"), MsgValue::Binary(data)),
364                    (MsgValue::from("format"), optional_str(format)),
365                ]),
366            ));
367        }
368        if let Some(document) = &self.document {
369            let (data, format) = document.resolve()?;
370            fields.push((
371                MsgValue::from("document"),
372                MsgValue::Map(vec![
373                    (MsgValue::from("data"), MsgValue::Binary(data)),
374                    (MsgValue::from("format"), optional_str(format)),
375                ]),
376            ));
377        }
378        if let Some(metadata) = &self.metadata {
379            fields.push((MsgValue::from("metadata"), json_to_msgpack(metadata)));
380        }
381        Ok(MsgValue::Map(fields))
382    }
383}
384
385fn optional_str(value: Option<String>) -> MsgValue {
386    value.map_or(MsgValue::Nil, MsgValue::from)
387}
388
389/// Translate a JSON value into its msgpack equivalent.
390pub(crate) fn json_to_msgpack(value: &Value) -> MsgValue {
391    match value {
392        Value::Null => MsgValue::Nil,
393        Value::Bool(flag) => MsgValue::Boolean(*flag),
394        Value::Number(number) => number.as_i64().map_or_else(
395            || {
396                number.as_u64().map_or_else(
397                    || MsgValue::from(number.as_f64().unwrap_or(0.0)),
398                    MsgValue::from,
399                )
400            },
401            MsgValue::from,
402        ),
403        Value::String(text) => MsgValue::from(text.as_str()),
404        Value::Array(items) => MsgValue::Array(items.iter().map(json_to_msgpack).collect()),
405        Value::Object(entries) => MsgValue::Map(
406            entries
407                .iter()
408                .map(|(key, value)| (MsgValue::from(key.as_str()), json_to_msgpack(value)))
409                .collect(),
410        ),
411    }
412}
413
414#[cfg(test)]
415mod tests {
416    // These assertions are about exact values, so exact comparison is the point.
417    #![allow(clippy::float_cmp)]
418
419    use super::*;
420    use serde_json::json;
421
422    fn field<'a>(value: &'a MsgValue, name: &str) -> Option<&'a MsgValue> {
423        match value {
424            MsgValue::Map(entries) => entries
425                .iter()
426                .find(|(key, _)| key.as_str() == Some(name))
427                .map(|(_, value)| value),
428            _ => None,
429        }
430    }
431
432    fn png_bytes() -> Vec<u8> {
433        let image = image::DynamicImage::ImageRgb8(image::RgbImage::new(2, 2));
434        let mut buffer = std::io::Cursor::new(Vec::new());
435        image
436            .write_to(&mut buffer, image::ImageFormat::Png)
437            .unwrap();
438        buffer.into_inner()
439    }
440
441    #[test]
442    fn a_text_item_emits_only_text() {
443        let wire = Item::text("hello").to_msgpack().unwrap();
444        assert_eq!(field(&wire, "text").unwrap().as_str(), Some("hello"));
445        assert!(field(&wire, "images").is_none());
446        assert!(field(&wire, "id").is_none());
447    }
448
449    #[test]
450    fn encoded_images_pass_through_with_a_detected_format() {
451        let png = png_bytes();
452        let wire = Item::image(ImageInput::bytes(png.clone()))
453            .to_msgpack()
454            .unwrap();
455        let images = field(&wire, "images").unwrap();
456        let MsgValue::Array(images) = images else {
457            panic!("expected an array")
458        };
459        assert_eq!(field(&images[0], "format").unwrap().as_str(), Some("png"));
460        // The bytes are carried verbatim: an already-encoded image is never re-encoded.
461        assert_eq!(
462            field(&images[0], "data").unwrap(),
463            &MsgValue::Binary(png),
464            "image bytes were rewritten"
465        );
466    }
467
468    #[test]
469    fn decoded_images_become_jpeg() {
470        let image = image::DynamicImage::ImageRgb8(image::RgbImage::new(2, 2));
471        let (data, format) = ImageInput::decoded(image).resolve().unwrap();
472        assert_eq!(format, "jpeg");
473        assert_eq!(media::detect_image_format(&data).unwrap(), "jpeg");
474    }
475
476    #[test]
477    fn a_declared_format_that_contradicts_the_bytes_is_rejected() {
478        let err = ImageInput::bytes(png_bytes())
479            .format("jpeg")
480            .resolve()
481            .unwrap_err();
482        assert!(err.to_string().contains("mismatch"), "{err}");
483        // The alias still matches its canonical form.
484        assert!(
485            ImageInput::bytes(png_bytes())
486                .format("PNG")
487                .resolve()
488                .is_ok()
489        );
490    }
491
492    #[test]
493    fn waveforms_are_wrapped_in_wav() {
494        let audio = AudioInput::waveform(Samples::F32(vec![0.0, 0.5, -0.5, 0.0]), 1, 16_000);
495        let (data, format, sample_rate) = audio.resolve().unwrap();
496        assert_eq!(format.as_deref(), Some("wav"));
497        assert_eq!(sample_rate, Some(16_000));
498        assert_eq!(&data[..4], b"RIFF");
499    }
500
501    #[test]
502    fn document_format_is_inferred_from_the_path_but_never_overrides_a_declaration() {
503        let dir = std::env::temp_dir().join("sie-sdk-item-tests");
504        std::fs::create_dir_all(&dir).unwrap();
505        let path = dir.join("report.pdf");
506        std::fs::write(&path, b"%PDF-1.4").unwrap();
507
508        let (data, format) = BinaryInput::document_path(&path).resolve().unwrap();
509        assert_eq!(data, b"%PDF-1.4");
510        assert_eq!(format.as_deref(), Some("pdf"));
511
512        let (_, declared) = BinaryInput::document_path(&path)
513            .format("txt")
514            .resolve()
515            .unwrap();
516        assert_eq!(declared.as_deref(), Some("txt"));
517
518        // In-memory documents carry no format unless the caller declares one.
519        assert_eq!(
520            BinaryInput::document_bytes(b"raw".to_vec())
521                .resolve()
522                .unwrap()
523                .1,
524            None
525        );
526        std::fs::remove_file(&path).unwrap();
527    }
528
529    #[test]
530    fn a_missing_file_names_itself_in_the_error() {
531        let err = ImageInput::path("/nonexistent/sie-sdk/image.png")
532            .resolve()
533            .unwrap_err();
534        assert!(
535            err.to_string().contains("/nonexistent/sie-sdk/image.png"),
536            "{err}"
537        );
538    }
539
540    #[test]
541    fn metadata_translates_into_msgpack() {
542        let wire = Item::text("x")
543            .with_id("doc-1")
544            .with_metadata(
545                json!({"source": "web", "rank": 3, "tags": ["a"], "keep": null, "score": 1.5}),
546            )
547            .to_msgpack()
548            .unwrap();
549        assert_eq!(field(&wire, "id").unwrap().as_str(), Some("doc-1"));
550        let metadata = field(&wire, "metadata").unwrap();
551        assert_eq!(field(metadata, "source").unwrap().as_str(), Some("web"));
552        assert_eq!(field(metadata, "rank").unwrap().as_i64(), Some(3));
553        assert_eq!(field(metadata, "score").unwrap().as_f64(), Some(1.5));
554        assert_eq!(field(metadata, "keep").unwrap(), &MsgValue::Nil);
555    }
556}
sie_sdk/types/item.rs

sie_sdk/types/
item.rs
