sie_sdk/types/
common.rs

1//! Scalar wire enums and the per-request metadata every endpoint can return.
2
3// Wire-mirror types: field names are the API contract itself, and the ones whose
4// meaning is not obvious carry their own doc comment.
5#![allow(missing_docs)]
6
7use serde::{Deserialize, Serialize};
8
9/// Compute dtype of a tensor on the wire.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum DType {
13    Float32,
14    Float16,
15    BFloat16,
16    Int8,
17    UInt8,
18    Binary,
19    UBinary,
20}
21
22/// Dtype a caller may request for returned embeddings.
23///
24/// `bfloat16` is compute-only and is never returned, so it has no variant here.
25#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
26#[serde(rename_all = "lowercase")]
27pub enum OutputDType {
28    #[default]
29    Float32,
30    Float16,
31    Int8,
32    UInt8,
33    Binary,
34    UBinary,
35}
36
37/// Which representations an encode call should return.
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
39#[serde(rename_all = "lowercase")]
40pub enum OutputType {
41    Dense,
42    Sparse,
43    Multivector,
44}
45
46/// Lifecycle state of a model on a worker.
47///
48/// Pinned to `packages/wire-fixtures/model_state.json` in the SIE repo, which is the
49/// cross-language source of truth.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
51#[serde(rename_all = "lowercase")]
52pub enum ModelState {
53    Available,
54    Loading,
55    Loaded,
56    Unloading,
57    Failed,
58}
59
60impl ModelState {
61    /// Every variant, in fixture order.
62    pub const ALL: [Self; 5] = [
63        Self::Available,
64        Self::Loading,
65        Self::Loaded,
66        Self::Unloading,
67        Self::Failed,
68    ];
69
70    /// The wire string for this state.
71    pub fn as_str(self) -> &'static str {
72        match self {
73            Self::Available => "available",
74            Self::Loading => "loading",
75            Self::Loaded => "loaded",
76            Self::Unloading => "unloading",
77            Self::Failed => "failed",
78        }
79    }
80}
81
82/// Billable units consumed by one request.
83///
84/// Populated from the `X-SIE-Units-*` response headers, or from a response body `usage`
85/// object when the server settled the charge inline.
86#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
87pub struct RequestUsage {
88    #[serde(default, skip_serializing_if = "Option::is_none")]
89    pub input_tokens: Option<u64>,
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub pairs: Option<u64>,
92    #[serde(default, skip_serializing_if = "Option::is_none")]
93    pub images: Option<u64>,
94    #[serde(default, skip_serializing_if = "Option::is_none")]
95    pub pages: Option<u64>,
96    #[serde(default, skip_serializing_if = "Option::is_none")]
97    pub output_tokens: Option<u64>,
98    #[serde(default, skip_serializing_if = "Option::is_none")]
99    pub audio_ms: Option<u64>,
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub credits_charged: Option<u64>,
102    #[serde(default, skip_serializing_if = "Option::is_none")]
103    pub rate_book_version: Option<String>,
104}
105
106impl RequestUsage {
107    pub(crate) fn is_empty(&self) -> bool {
108        *self == Self::default()
109    }
110}
111
112/// Server-reported metadata about a single request.
113///
114/// `retries` and `model_revision` are client-side observations; the rest come from the
115/// server's response headers or body.
116#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
117pub struct RequestMetadata {
118    #[serde(default, skip_serializing_if = "Option::is_none")]
119    pub id: Option<String>,
120    #[serde(default, skip_serializing_if = "Option::is_none")]
121    pub usage: Option<RequestUsage>,
122    #[serde(default, skip_serializing_if = "Option::is_none")]
123    pub credits_debited: Option<u64>,
124    #[serde(default, skip_serializing_if = "Option::is_none")]
125    pub rate_book_version: Option<String>,
126    #[serde(default, skip_serializing_if = "Option::is_none")]
127    pub execution_identity_sha256: Option<String>,
128    /// How many times the SDK retried before this response arrived.
129    #[serde(default)]
130    pub retries: u32,
131    /// Value of `X-SIE-Model-Revision`, when the server sent one.
132    #[serde(default, skip_serializing_if = "Option::is_none")]
133    pub model_revision: Option<String>,
134}
135
136impl RequestMetadata {
137    pub(crate) fn is_empty(&self) -> bool {
138        self.id.is_none()
139            && self.usage.is_none()
140            && self.credits_debited.is_none()
141            && self.rate_book_version.is_none()
142            && self.execution_identity_sha256.is_none()
143            && self.model_revision.is_none()
144    }
145}
146
147/// Per-stage latency breakdown returned by encode.
148#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
149pub struct TimingInfo {
150    #[serde(default, skip_serializing_if = "Option::is_none")]
151    pub total_ms: Option<f64>,
152    #[serde(default, skip_serializing_if = "Option::is_none")]
153    pub queue_ms: Option<f64>,
154    #[serde(default, skip_serializing_if = "Option::is_none")]
155    pub tokenization_ms: Option<f64>,
156    #[serde(default, skip_serializing_if = "Option::is_none")]
157    pub inference_ms: Option<f64>,
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    #[test]
165    fn wire_strings_round_trip() {
166        assert_eq!(
167            serde_json::to_string(&DType::BFloat16).unwrap(),
168            "\"bfloat16\""
169        );
170        assert_eq!(
171            serde_json::to_string(&OutputType::Multivector).unwrap(),
172            "\"multivector\""
173        );
174        assert_eq!(
175            serde_json::to_string(&OutputDType::UBinary).unwrap(),
176            "\"ubinary\""
177        );
178        for state in ModelState::ALL {
179            let json = serde_json::to_string(&state).unwrap();
180            assert_eq!(json, format!("\"{}\"", state.as_str()));
181            assert_eq!(serde_json::from_str::<ModelState>(&json).unwrap(), state);
182        }
183    }
184}
sie_sdk/types/common.rs

sie_sdk/types/
common.rs
