sie_sdk/client/
encode.rs

1//! `/v1/encode`, `/v1/score` and `/v1/extract`: the msgpack inference endpoints.
2
3use reqwest::Method;
4use rmpv::Value as MsgValue;
5use serde_json::Value;
6
7use crate::client::Client;
8use crate::error::{Error, Result, codes};
9use crate::http::{HttpResponse, headers, metadata};
10use crate::retry::{RequestOptions, RetryPolicy, RetryState};
11use crate::types::item::json_to_msgpack;
12use crate::types::{
13    Classification, DetectedObject, EncodeResult, Entity, ExtractItemError, ExtractResult, Item,
14    Multivector, OutputDType, OutputType, Relation, RequestMetadata, ScoreEntry, ScoreResult,
15    ScoreUsage, SparseVector, TimingInfo,
16};
17use crate::wire::{msg, ndarray};
18
19const MALFORMED_EXTRACT_ERROR_MESSAGE: &str = "Malformed extraction item error";
20
21/// Options shared by every request builder in this module.
22macro_rules! request_options {
23    () => {
24        /// Route to a machine profile, optionally pool-qualified as `"pool/profile"`.
25        pub fn gpu(mut self, gpu: impl Into<String>) -> Self {
26            self.options.gpu = Some(gpu.into());
27            self
28        }
29
30        /// Whether to wait out provisioning and model loading, or fail fast.
31        pub fn wait_for_capacity(mut self, wait: bool) -> Self {
32            self.options.wait_for_capacity = wait;
33            self
34        }
35
36        /// Total wall-clock budget for this call, retries included.
37        pub fn provision_timeout(mut self, timeout: std::time::Duration) -> Self {
38            self.options.provision_timeout = timeout;
39            self
40        }
41
42        /// Cap on `RESOURCE_EXHAUSTED` retries. Zero fails fast.
43        pub fn max_oom_retries(mut self, retries: u32) -> Self {
44            self.options.max_oom_retries = retries;
45            self
46        }
47    };
48}
49
50pub(crate) use request_options;
51
52impl Client {
53    /// Embed one or more items.
54    pub fn encode(
55        &self,
56        model: impl Into<String>,
57        items: impl IntoIterator<Item = Item>,
58    ) -> EncodeRequest {
59        EncodeRequest {
60            client: self.clone(),
61            model: model.into(),
62            items: items.into_iter().collect(),
63            output_types: None,
64            instruction: None,
65            output_dtype: None,
66            is_query: None,
67            runtime_options: None,
68            options: self.request_options(),
69        }
70    }
71
72    /// Score a query against a set of candidates.
73    pub fn score(
74        &self,
75        model: impl Into<String>,
76        query: Item,
77        items: impl IntoIterator<Item = Item>,
78    ) -> ScoreRequest {
79        ScoreRequest {
80            client: self.clone(),
81            model: model.into(),
82            query,
83            items: items.into_iter().collect(),
84            instruction: None,
85            runtime_options: None,
86            options: self.request_options(),
87        }
88    }
89
90    /// Extract structure from one or more items.
91    pub fn extract(
92        &self,
93        model: impl Into<String>,
94        items: impl IntoIterator<Item = Item>,
95    ) -> ExtractRequest {
96        ExtractRequest {
97            client: self.clone(),
98            model: model.into(),
99            items: items.into_iter().collect(),
100            labels: None,
101            output_schema: None,
102            instruction: None,
103            runtime_options: None,
104            options: self.request_options(),
105        }
106    }
107
108    /// Send a msgpack body and decode the response envelope.
109    async fn send_msgpack(
110        &self,
111        path: &str,
112        body: MsgValue,
113        policy: RetryPolicy,
114        model: &str,
115        options: &RequestOptions,
116    ) -> Result<(MsgValue, HttpResponse, u32)> {
117        let routing = self.routing(options.gpu.as_deref());
118        let encoded = rmp_serde::to_vec(&body)
119            .map_err(|err| Error::invalid(format!("could not encode the request body: {err}")))?;
120
121        let request = self
122            .request(Method::POST, path)?
123            .msgpack_headers()
124            .maybe_header(headers::MACHINE_PROFILE, routing.profile.as_deref())
125            .maybe_header(headers::POOL, routing.pool.as_deref())
126            .body(encoded);
127
128        let mut state = RetryState::new(policy, options, Some(model));
129        let response = self.send(request, &mut state).await?;
130        let decoded: MsgValue = rmp_serde::from_slice(&response.body)
131            .map_err(|err| Error::decode(format!("malformed msgpack response: {err}")))?;
132        let retries = state.retries();
133        Ok((decoded, response, retries))
134    }
135}
136
137/// Assemble the `params` sub-map, omitting it entirely when nothing was set.
138fn params_map(entries: Vec<(&str, MsgValue)>) -> Option<MsgValue> {
139    if entries.is_empty() {
140        return None;
141    }
142    Some(MsgValue::Map(
143        entries
144            .into_iter()
145            .map(|(key, value)| (MsgValue::from(key), value))
146            .collect(),
147    ))
148}
149
150fn items_to_msgpack(items: &[Item]) -> Result<MsgValue> {
151    Ok(MsgValue::Array(
152        items
153            .iter()
154            .map(Item::to_msgpack)
155            .collect::<Result<Vec<_>>>()?,
156    ))
157}
158
159/// The response envelope's own `usage` block, for the settled charge.
160fn body_usage(body: &MsgValue) -> Option<Value> {
161    Some(msg::to_json(&MsgValue::Map(vec![(
162        MsgValue::from("usage"),
163        msg::get(body, "usage")?.clone(),
164    )])))
165}
166
167fn attach_metadata(
168    response: &HttpResponse,
169    body: &MsgValue,
170    retries: u32,
171) -> Option<RequestMetadata> {
172    metadata::parse(&response.headers, body_usage(body).as_ref(), retries)
173}
174
175fn timing(body: &MsgValue) -> Option<TimingInfo> {
176    let timing = msg::get(body, "timing")?;
177    Some(TimingInfo {
178        total_ms: msg::get_f64(timing, "total_ms"),
179        queue_ms: msg::get_f64(timing, "queue_ms"),
180        tokenization_ms: msg::get_f64(timing, "tokenization_ms"),
181        inference_ms: msg::get_f64(timing, "inference_ms"),
182    })
183}
184
185/// Embeds one or more items. Build with [`Client::encode`].
186pub struct EncodeRequest {
187    client: Client,
188    model: String,
189    items: Vec<Item>,
190    output_types: Option<Vec<OutputType>>,
191    instruction: Option<String>,
192    output_dtype: Option<OutputDType>,
193    is_query: Option<bool>,
194    runtime_options: Option<Value>,
195    options: RequestOptions,
196}
197
198impl EncodeRequest {
199    request_options!();
200
201    /// Which representations to return. Defaults to whatever the model produces.
202    pub fn output_types(mut self, types: impl IntoIterator<Item = OutputType>) -> Self {
203        self.output_types = Some(types.into_iter().collect());
204        self
205    }
206
207    /// Task instruction prepended by instruction-tuned models.
208    pub fn instruction(mut self, instruction: impl Into<String>) -> Self {
209        self.instruction = Some(instruction.into());
210        self
211    }
212
213    /// Quantize the returned embeddings.
214    pub fn output_dtype(mut self, dtype: OutputDType) -> Self {
215        self.output_dtype = Some(dtype);
216        self
217    }
218
219    /// Encode as a query rather than a document, for asymmetric models.
220    pub fn is_query(mut self, is_query: bool) -> Self {
221        self.is_query = Some(is_query);
222        self
223    }
224
225    /// Adapter-specific runtime options, merged over the client defaults.
226    pub fn options(mut self, options: Value) -> Self {
227        self.runtime_options = Some(options);
228        self
229    }
230
231    fn body(&self) -> Result<MsgValue> {
232        let mut options = self.client.merge_options(self.runtime_options.as_ref());
233        // `is_query` rides inside `options`, not alongside it.
234        if let Some(is_query) = self.is_query {
235            let mut map = options
236                .take()
237                .and_then(|value| value.as_object().cloned())
238                .unwrap_or_default();
239            map.insert("is_query".to_string(), Value::Bool(is_query));
240            options = Some(Value::Object(map));
241        }
242
243        let mut params = Vec::new();
244        if let Some(types) = &self.output_types {
245            params.push((
246                "output_types",
247                MsgValue::Array(
248                    types
249                        .iter()
250                        .map(|kind| {
251                            MsgValue::from(
252                                serde_json::to_value(kind)
253                                    .unwrap_or_default()
254                                    .as_str()
255                                    .unwrap_or(""),
256                            )
257                        })
258                        .collect(),
259                ),
260            ));
261        }
262        if let Some(instruction) = &self.instruction {
263            params.push(("instruction", MsgValue::from(instruction.as_str())));
264        }
265        if let Some(dtype) = self.output_dtype {
266            params.push((
267                "output_dtype",
268                MsgValue::from(
269                    serde_json::to_value(dtype)
270                        .unwrap_or_default()
271                        .as_str()
272                        .unwrap_or(""),
273                ),
274            ));
275        }
276        if let Some(options) = &options {
277            params.push(("options", json_to_msgpack(options)));
278        }
279
280        let mut fields = vec![(MsgValue::from("items"), items_to_msgpack(&self.items)?)];
281        if let Some(params) = params_map(params) {
282            fields.push((MsgValue::from("params"), params));
283        }
284        Ok(MsgValue::Map(fields))
285    }
286
287    /// Send the request.
288    pub async fn send(self) -> Result<Vec<EncodeResult>> {
289        if self.items.is_empty() {
290            return Err(Error::invalid("encode requires at least one item"));
291        }
292        let expected = self.items.len();
293        let body = self.body()?;
294        let (decoded, response, retries) = self
295            .client
296            .send_msgpack(
297                &format!("/v1/encode/{}", self.model),
298                body,
299                RetryPolicy::ENCODE,
300                &self.model,
301                &self.options,
302            )
303            .await?;
304
305        let items = msg::get_array(&decoded, "items")
306            .ok_or_else(|| Error::decode("encode response is missing its `items` array"))?;
307        // The 1:1 input/output contract is load-bearing: callers reassemble batches
308        // positionally, so a desynced response must never reach them.
309        if items.len() != expected {
310            return Err(Error::Server {
311                message: format!(
312                    "Server returned {} results for {expected} items (model '{}')",
313                    items.len(),
314                    self.model
315                ),
316                code: Some(codes::ENCODE_RESULT_COUNT_MISMATCH.to_string()),
317                status: response.status,
318                request: attach_metadata(&response, &decoded, retries).map(Box::new),
319            });
320        }
321
322        let model = msg::get_string(&decoded, "model");
323        let timing = timing(&decoded);
324        let request = attach_metadata(&response, &decoded, retries);
325
326        items
327            .iter()
328            .map(|item| {
329                Ok(EncodeResult {
330                    model: model.clone(),
331                    id: msg::get_string(item, "id"),
332                    dense: parse_dense(item)?,
333                    sparse: parse_sparse(item)?,
334                    multivector: parse_multivector(item)?,
335                    timing,
336                    request: request.clone(),
337                })
338            })
339            .collect()
340    }
341
342    /// Send the request and return the single expected result.
343    pub async fn send_one(self) -> Result<EncodeResult> {
344        let mut results = self.send().await?;
345        match results.len() {
346            1 => Ok(results.remove(0)),
347            other => Err(Error::invalid(format!(
348                "send_one() expects exactly one item, but the request carried {other}"
349            ))),
350        }
351    }
352}
353
354/// The wire nests each tensor under a `values` key; the SDK type does not.
355fn tensor(item: &MsgValue, key: &str) -> Result<Option<ndarray::RawArray>> {
356    let Some(node) = msg::get(item, key) else {
357        return Ok(None);
358    };
359    let values = msg::get(node, "values").ok_or_else(|| {
360        Error::decode(format!(
361            "encode result `{key}` is missing its `values` array"
362        ))
363    })?;
364    Ok(Some(ndarray::decode(values)?))
365}
366
367fn parse_dense(item: &MsgValue) -> Result<Option<Vec<f32>>> {
368    Ok(tensor(item, "dense")?.map(|array| array.to_f32()))
369}
370
371fn parse_multivector(item: &MsgValue) -> Result<Option<Multivector>> {
372    let Some(array) = tensor(item, "multivector")? else {
373        return Ok(None);
374    };
375    Ok(Some(match array.element {
376        ndarray::Element::F16 => Multivector::F16(array.rows_f16()?),
377        _ => Multivector::F32(array.rows_f32()?),
378    }))
379}
380
381fn parse_sparse(item: &MsgValue) -> Result<Option<SparseVector>> {
382    let Some(node) = msg::get(item, "sparse") else {
383        return Ok(None);
384    };
385    let indices = msg::get(node, "indices")
386        .ok_or_else(|| Error::decode("encode result `sparse` is missing its `indices` array"))?;
387    let values = msg::get(node, "values")
388        .ok_or_else(|| Error::decode("encode result `sparse` is missing its `values` array"))?;
389    Ok(Some(SparseVector {
390        indices: ndarray::decode(indices)?.to_u32()?,
391        values: ndarray::decode(values)?.to_f32(),
392    }))
393}
394
395/// Scores a query against candidates. Build with [`Client::score`].
396pub struct ScoreRequest {
397    client: Client,
398    model: String,
399    query: Item,
400    items: Vec<Item>,
401    instruction: Option<String>,
402    runtime_options: Option<Value>,
403    options: RequestOptions,
404}
405
406impl ScoreRequest {
407    request_options!();
408
409    /// Task instruction prepended by instruction-tuned rerankers.
410    pub fn instruction(mut self, instruction: impl Into<String>) -> Self {
411        self.instruction = Some(instruction.into());
412        self
413    }
414
415    /// Adapter-specific runtime options, merged over the client defaults.
416    pub fn options(mut self, options: Value) -> Self {
417        self.runtime_options = Some(options);
418        self
419    }
420
421    /// Send the request.
422    pub async fn send(self) -> Result<ScoreResult> {
423        if self.items.is_empty() {
424            return Err(Error::invalid("score requires at least one candidate item"));
425        }
426
427        // Score fields sit at the top level of the body, not inside `params`.
428        let mut fields = vec![
429            (MsgValue::from("query"), self.query.to_msgpack()?),
430            (MsgValue::from("items"), items_to_msgpack(&self.items)?),
431        ];
432        if let Some(instruction) = &self.instruction {
433            fields.push((
434                MsgValue::from("instruction"),
435                MsgValue::from(instruction.as_str()),
436            ));
437        }
438        if let Some(options) = self.client.merge_options(self.runtime_options.as_ref()) {
439            fields.push((MsgValue::from("options"), json_to_msgpack(&options)));
440        }
441
442        let (decoded, response, retries) = self
443            .client
444            .send_msgpack(
445                &format!("/v1/score/{}", self.model),
446                MsgValue::Map(fields),
447                RetryPolicy::INFERENCE,
448                &self.model,
449                &self.options,
450            )
451            .await?;
452
453        let model = msg::get_string(&decoded, "model")
454            .ok_or_else(|| Error::decode("score response is missing its `model` field"))?;
455        let entries = msg::get_array(&decoded, "scores")
456            .ok_or_else(|| Error::decode("score response is missing its `scores` array"))?;
457
458        let scores = entries
459            .iter()
460            .map(|entry| {
461                Ok(ScoreEntry {
462                    item_id: msg::get_string(entry, "item_id")
463                        .ok_or_else(|| Error::decode("score entry is missing its `item_id`"))?,
464                    score: msg::get_f64(entry, "score")
465                        .ok_or_else(|| Error::decode("score entry is missing its `score`"))?,
466                    rank: msg::get_u64(entry, "rank").unwrap_or(0) as u32,
467                })
468            })
469            .collect::<Result<Vec<_>>>()?;
470
471        Ok(ScoreResult {
472            model,
473            query_id: msg::get_string(&decoded, "query_id"),
474            scores,
475            usage: msg::get(&decoded, "usage").map(|usage| ScoreUsage {
476                input_tokens: msg::get_u64(usage, "input_tokens").unwrap_or(0),
477                images: msg::get_u64(usage, "images"),
478            }),
479            request: attach_metadata(&response, &decoded, retries),
480        })
481    }
482}
483
484/// Extracts structure from items. Build with [`Client::extract`].
485pub struct ExtractRequest {
486    client: Client,
487    model: String,
488    items: Vec<Item>,
489    labels: Option<Vec<String>>,
490    output_schema: Option<Value>,
491    instruction: Option<String>,
492    runtime_options: Option<Value>,
493    options: RequestOptions,
494}
495
496impl ExtractRequest {
497    request_options!();
498
499    /// Candidate labels, for zero-shot classifiers and detectors.
500    pub fn labels(mut self, labels: impl IntoIterator<Item = impl Into<String>>) -> Self {
501        self.labels = Some(labels.into_iter().map(Into::into).collect());
502        self
503    }
504
505    /// JSON schema the model should populate.
506    pub fn output_schema(mut self, schema: Value) -> Self {
507        self.output_schema = Some(schema);
508        self
509    }
510
511    /// Task instruction.
512    pub fn instruction(mut self, instruction: impl Into<String>) -> Self {
513        self.instruction = Some(instruction.into());
514        self
515    }
516
517    /// Adapter-specific runtime options, merged over the client defaults.
518    pub fn options(mut self, options: Value) -> Self {
519        self.runtime_options = Some(options);
520        self
521    }
522
523    /// Send the request.
524    pub async fn send(self) -> Result<Vec<ExtractResult>> {
525        if self.items.is_empty() {
526            return Err(Error::invalid("extract requires at least one item"));
527        }
528
529        let mut params = Vec::new();
530        if let Some(labels) = &self.labels {
531            params.push((
532                "labels",
533                MsgValue::Array(
534                    labels
535                        .iter()
536                        .map(|label| MsgValue::from(label.as_str()))
537                        .collect(),
538                ),
539            ));
540        }
541        if let Some(schema) = &self.output_schema {
542            params.push(("output_schema", json_to_msgpack(schema)));
543        }
544        if let Some(instruction) = &self.instruction {
545            params.push(("instruction", MsgValue::from(instruction.as_str())));
546        }
547        if let Some(options) = self.client.merge_options(self.runtime_options.as_ref()) {
548            params.push(("options", json_to_msgpack(&options)));
549        }
550
551        let mut fields = vec![(MsgValue::from("items"), items_to_msgpack(&self.items)?)];
552        if let Some(params) = params_map(params) {
553            fields.push((MsgValue::from("params"), params));
554        }
555
556        let (decoded, response, retries) = self
557            .client
558            .send_msgpack(
559                &format!("/v1/extract/{}", self.model),
560                MsgValue::Map(fields),
561                RetryPolicy::INFERENCE,
562                &self.model,
563                &self.options,
564            )
565            .await?;
566
567        let items = msg::get_array(&decoded, "items")
568            .ok_or_else(|| Error::decode("extract response is missing its `items` array"))?;
569        let model = msg::get_string(&decoded, "model");
570        let request = attach_metadata(&response, &decoded, retries);
571
572        Ok(items
573            .iter()
574            .map(|item| ExtractResult {
575                model: model.clone(),
576                id: msg::get_string(item, "id"),
577                entities: parse_list(item, "entities", parse_entity),
578                relations: parse_list(item, "relations", parse_relation),
579                classifications: parse_list(item, "classifications", parse_classification),
580                objects: parse_list(item, "objects", parse_object),
581                data: msg::get(item, "data")
582                    .map(msg::to_json)
583                    .filter(|data| !is_blank(data)),
584                error: msg::get(item, "error").map(parse_extract_error),
585                request: request.clone(),
586            })
587            .collect())
588    }
589
590    /// Send the request and return the single expected result.
591    pub async fn send_one(self) -> Result<ExtractResult> {
592        let mut results = self.send().await?;
593        match results.len() {
594            1 => Ok(results.remove(0)),
595            other => Err(Error::invalid(format!(
596                "send_one() expects exactly one item, but the response carried {other}"
597            ))),
598        }
599    }
600}
601
602fn is_blank(value: &Value) -> bool {
603    match value {
604        Value::Null => true,
605        Value::Object(map) => map.is_empty(),
606        Value::Array(items) => items.is_empty(),
607        _ => false,
608    }
609}
610
611fn parse_list<T>(item: &MsgValue, key: &str, parse: fn(&MsgValue) -> T) -> Vec<T> {
612    msg::get_array(item, key)
613        .map(|entries| entries.iter().map(parse).collect())
614        .unwrap_or_default()
615}
616
617fn parse_entity(value: &MsgValue) -> Entity {
618    Entity {
619        text: msg::get_string(value, "text").unwrap_or_default(),
620        label: msg::get_string(value, "label").unwrap_or_default(),
621        score: msg::get_f64(value, "score").unwrap_or_default(),
622        start: msg::get_i64(value, "start"),
623        end: msg::get_i64(value, "end"),
624        bbox: parse_bbox(value),
625    }
626}
627
628fn parse_relation(value: &MsgValue) -> Relation {
629    Relation {
630        head: msg::get_string(value, "head").unwrap_or_default(),
631        tail: msg::get_string(value, "tail").unwrap_or_default(),
632        relation: msg::get_string(value, "relation").unwrap_or_default(),
633        score: msg::get_f64(value, "score").unwrap_or_default(),
634    }
635}
636
637fn parse_classification(value: &MsgValue) -> Classification {
638    Classification {
639        label: msg::get_string(value, "label").unwrap_or_default(),
640        score: msg::get_f64(value, "score").unwrap_or_default(),
641    }
642}
643
644fn parse_object(value: &MsgValue) -> DetectedObject {
645    DetectedObject {
646        label: msg::get_string(value, "label").unwrap_or_default(),
647        score: msg::get_f64(value, "score").unwrap_or_default(),
648        bbox: parse_bbox(value).unwrap_or_default(),
649    }
650}
651
652fn parse_bbox(value: &MsgValue) -> Option<Vec<i64>> {
653    Some(
654        msg::get_array(value, "bbox")?
655            .iter()
656            .filter_map(|dim| dim.as_i64().or_else(|| dim.as_f64().map(|v| v as i64)))
657            .collect(),
658    )
659}
660
661/// A per-item error is preserved only when both halves are present and non-blank;
662/// anything else is normalized so callers never see a half-formed failure.
663fn parse_extract_error(value: &MsgValue) -> ExtractItemError {
664    let code = msg::get_text(value, "code")
665        .map(str::trim)
666        .unwrap_or_default();
667    let message = msg::get_text(value, "message")
668        .map(str::trim)
669        .unwrap_or_default();
670    if code.is_empty() || message.is_empty() {
671        return ExtractItemError {
672            code: codes::INTERNAL_ERROR.to_string(),
673            message: MALFORMED_EXTRACT_ERROR_MESSAGE.to_string(),
674        };
675    }
676    ExtractItemError {
677        code: code.to_string(),
678        message: message.to_string(),
679    }
680}
681
682#[cfg(test)]
683mod tests {
684    // These assertions are about exact values, so exact comparison is the point.
685    #![allow(clippy::float_cmp)]
686
687    use super::*;
688    use crate::wire::ndarray::fixtures;
689
690    fn map(entries: Vec<(&str, MsgValue)>) -> MsgValue {
691        MsgValue::Map(
692            entries
693                .into_iter()
694                .map(|(key, value)| (MsgValue::from(key), value))
695                .collect(),
696        )
697    }
698
699    #[test]
700    fn unwraps_the_nested_tensor_envelopes() {
701        let item = map(vec![
702            ("id", MsgValue::from("doc-1")),
703            (
704                "dense",
705                map(vec![(
706                    "values",
707                    fixtures::f32_array(&[3], &[0.1, 0.2, 0.3]),
708                )]),
709            ),
710            (
711                "sparse",
712                map(vec![
713                    ("indices", fixtures::i32_array(&[2], &[7, 42])),
714                    ("values", fixtures::f32_array(&[2], &[0.5, 0.25])),
715                ]),
716            ),
717            (
718                "multivector",
719                map(vec![(
720                    "values",
721                    fixtures::f16_array(&[2, 2], &[1.0, 0.0, 0.0, 1.0]),
722                )]),
723            ),
724        ]);
725
726        assert_eq!(parse_dense(&item).unwrap().unwrap(), vec![0.1, 0.2, 0.3]);
727        let sparse = parse_sparse(&item).unwrap().unwrap();
728        assert_eq!(sparse.indices, vec![7, 42]);
729        assert_eq!(sparse.values, vec![0.5, 0.25]);
730        let mv = parse_multivector(&item).unwrap().unwrap();
731        assert!(matches!(mv, Multivector::F16(_)));
732        assert_eq!(mv.to_f32(), vec![vec![1.0, 0.0], vec![0.0, 1.0]]);
733    }
734
735    #[test]
736    fn absent_tensors_stay_absent() {
737        let item = map(vec![("id", MsgValue::from("doc-1"))]);
738        assert!(parse_dense(&item).unwrap().is_none());
739        assert!(parse_sparse(&item).unwrap().is_none());
740        assert!(parse_multivector(&item).unwrap().is_none());
741    }
742
743    #[test]
744    fn a_tensor_without_its_values_key_is_a_decode_error() {
745        let item = map(vec![("dense", map(vec![("dims", MsgValue::from(3u64))]))]);
746        let err = parse_dense(&item).unwrap_err();
747        assert!(err.to_string().contains("values"), "{err}");
748    }
749
750    #[test]
751    fn well_formed_extract_errors_pass_through() {
752        let error = parse_extract_error(&map(vec![
753            ("code", MsgValue::from("INFERENCE_ERROR")),
754            ("message", MsgValue::from("Document export failed")),
755        ]));
756        assert_eq!(error.code, "INFERENCE_ERROR");
757        assert_eq!(error.message, "Document export failed");
758    }
759
760    #[test]
761    fn malformed_extract_errors_are_normalized() {
762        let cases = vec![
763            map(vec![("code", MsgValue::from("INFERENCE_ERROR"))]),
764            map(vec![("message", MsgValue::from("only a message"))]),
765            map(vec![
766                ("code", MsgValue::from(" ")),
767                ("message", MsgValue::from("\t")),
768            ]),
769            map(vec![]),
770        ];
771        for case in cases {
772            let error = parse_extract_error(&case);
773            assert_eq!(error.code, codes::INTERNAL_ERROR);
774            assert_eq!(error.message, MALFORMED_EXTRACT_ERROR_MESSAGE);
775        }
776    }
777
778    #[test]
779    fn extract_lists_default_to_empty_and_data_drops_when_blank() {
780        let item = map(vec![("data", MsgValue::Map(Vec::new()))]);
781        assert!(parse_list(&item, "entities", parse_entity).is_empty());
782        assert!(
783            msg::get(&item, "data")
784                .map(msg::to_json)
785                .as_ref()
786                .is_none_or(is_blank)
787        );
788    }
789
790    #[test]
791    fn entity_and_object_shapes_decode() {
792        let entity = parse_entity(&map(vec![
793            ("text", MsgValue::from("Ada")),
794            ("label", MsgValue::from("PERSON")),
795            ("score", MsgValue::F64(0.91)),
796            ("start", MsgValue::from(0i64)),
797            ("end", MsgValue::from(3i64)),
798        ]));
799        assert_eq!(entity.text, "Ada");
800        assert_eq!(entity.start, Some(0));
801        assert!(entity.bbox.is_none());
802
803        let object = parse_object(&map(vec![
804            ("label", MsgValue::from("cat")),
805            ("score", MsgValue::F64(0.7)),
806            (
807                "bbox",
808                MsgValue::Array(vec![
809                    MsgValue::from(1i64),
810                    MsgValue::from(2i64),
811                    MsgValue::from(30i64),
812                    MsgValue::from(40i64),
813                ]),
814            ),
815        ]));
816        assert_eq!(object.bbox, vec![1, 2, 30, 40]);
817    }
818
819    #[tokio::test]
820    async fn empty_input_is_rejected_before_any_request() {
821        let client = Client::new("https://sie.invalid").unwrap();
822        assert!(client.encode("m", Vec::new()).send().await.is_err());
823        assert!(client.extract("m", Vec::new()).send().await.is_err());
824        assert!(
825            client
826                .score("m", Item::text("q"), Vec::new())
827                .send()
828                .await
829                .is_err()
830        );
831    }
832
833    #[test]
834    fn encode_body_nests_params_and_folds_is_query_into_options() {
835        let client = Client::new("https://sie.invalid").unwrap();
836        let body = client
837            .encode("m", [Item::text("hello")])
838            .output_types([OutputType::Dense, OutputType::Sparse])
839            .instruction("query:")
840            .output_dtype(OutputDType::Int8)
841            .is_query(true)
842            .body()
843            .unwrap();
844
845        let params = msg::get(&body, "params").unwrap();
846        assert_eq!(msg::get_text(params, "instruction"), Some("query:"));
847        assert_eq!(msg::get_text(params, "output_dtype"), Some("int8"));
848        let types = msg::get_array(params, "output_types").unwrap();
849        assert_eq!(types[0].as_str(), Some("dense"));
850        assert_eq!(types[1].as_str(), Some("sparse"));
851        let options = msg::get(params, "options").unwrap();
852        assert_eq!(
853            msg::get(options, "is_query").unwrap(),
854            &MsgValue::Boolean(true)
855        );
856        assert!(msg::get_array(&body, "items").is_some());
857    }
858
859    #[test]
860    fn encode_body_omits_params_when_nothing_was_set() {
861        let client = Client::new("https://sie.invalid").unwrap();
862        let body = client.encode("m", [Item::text("hello")]).body().unwrap();
863        assert!(msg::get(&body, "params").is_none());
864    }
865}
sie_sdk/client/encode.rs

sie_sdk/client/
encode.rs
