rskit_inference/
types.rs

1use std::collections::HashMap;
2
3use bytes::Bytes;
4pub use rskit_ai::{
5    Capabilities, Model, Provider as ModelProvider, StreamEvent, StreamEventRef, Usage,
6};
7use rskit_errors::{AppError, ErrorCode};
8use serde::{Deserialize, Serialize};
9use thiserror::Error;
10
11/// Runtime-neutral model prediction request.
12#[derive(Debug, Clone, Default, Serialize, Deserialize)]
13pub struct PredictRequest {
14    /// Optional request id. Adapters generate UUID v7 when absent.
15    #[serde(default, skip_serializing_if = "Option::is_none")]
16    pub request_id: Option<String>,
17    /// Runtime model name.
18    pub model_name: String,
19    /// Optional runtime model version.
20    #[serde(default, skip_serializing_if = "Option::is_none")]
21    pub model_version: Option<String>,
22    /// Named inputs consumed by the serving runtime.
23    #[serde(default)]
24    pub inputs: HashMap<String, Value>,
25    /// Adapter-specific knobs such as `max_new_tokens`, `temperature`, or `top_k`.
26    #[serde(default)]
27    pub parameters: HashMap<String, serde_json::Value>,
28    /// Provider-specific runtime options.
29    #[serde(default)]
30    pub options: serde_json::Value,
31    /// Request metadata for authz, tracing, and adapter-specific headers.
32    #[serde(default)]
33    pub metadata: HashMap<String, String>,
34}
35
36/// Runtime-neutral prediction response.
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct PredictResponse {
39    /// Named outputs returned by the serving runtime.
40    #[serde(default)]
41    pub outputs: HashMap<String, Value>,
42    /// Token and compute usage counters.
43    #[serde(default)]
44    pub usage: Usage,
45    /// Model that served the request.
46    pub model: Model,
47    /// Prediction status.
48    pub status: PredictStatus,
49    /// Response metadata such as model version or finish reason.
50    #[serde(default)]
51    pub metadata: HashMap<String, String>,
52}
53
54impl Default for PredictResponse {
55    fn default() -> Self {
56        Self {
57            outputs: HashMap::new(),
58            usage: Usage::default(),
59            model: Model {
60                name: String::new(),
61                provider: ModelProvider::Custom("unknown".to_string()),
62                version: None,
63                capabilities: Capabilities::default(),
64            },
65            status: PredictStatus::Success,
66            metadata: HashMap::new(),
67        }
68    }
69}
70
71/// Prediction status.
72#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
73#[serde(rename_all = "snake_case")]
74#[non_exhaustive]
75pub enum PredictStatus {
76    /// Prediction completed successfully.
77    Success,
78    /// Prediction completed with partial outputs.
79    PartialSuccess,
80    /// Serving runtime reported an error.
81    Error {
82        /// Error reason.
83        reason: String,
84    },
85}
86
87/// Runtime-neutral input or output value.
88#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
89#[serde(tag = "kind", rename_all = "snake_case")]
90pub enum Value {
91    /// UTF-8 text value.
92    Text {
93        /// Text content.
94        text: String,
95    },
96    /// Opaque bytes value.
97    Bytes {
98        /// Byte content.
99        bytes: Bytes,
100    },
101    /// KServe v2-style tensor value.
102    Tensor {
103        /// Tensor content.
104        tensor: Tensor,
105    },
106    /// Arbitrary JSON value for runtimes that accept structured values directly.
107    Json {
108        /// JSON content.
109        json: serde_json::Value,
110    },
111}
112
113/// KServe v2-style tensor value.
114#[derive(Debug, Clone, PartialEq)]
115pub struct Tensor {
116    /// KServe v2 dtype name: `FP32`, `FP64`, `INT32`, `INT64`, `UINT8`, `BOOL`, `BYTES`, ...
117    pub dtype: String,
118    /// Tensor shape.
119    pub shape: Vec<i64>,
120    /// Tensor data.
121    pub data: TensorData,
122}
123
124impl Serialize for Tensor {
125    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
126    where
127        S: serde::Serializer,
128    {
129        use serde::ser::SerializeStruct;
130
131        let mut state = serializer.serialize_struct("Tensor", 3)?;
132        state.serialize_field("dtype", &self.dtype)?;
133        state.serialize_field("shape", &self.shape)?;
134        state.serialize_field("data", &self.data)?;
135        state.end()
136    }
137}
138
139impl<'de> Deserialize<'de> for Tensor {
140    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
141    where
142        D: serde::Deserializer<'de>,
143    {
144        #[derive(Deserialize)]
145        struct TensorWire {
146            dtype: String,
147            shape: Vec<i64>,
148            data: serde_json::Value,
149        }
150
151        let wire = TensorWire::deserialize(deserializer)?;
152        let dtype = wire.dtype.to_ascii_uppercase();
153        let data = match dtype.as_str() {
154            "FP32" => TensorData::F32(
155                serde_json::from_value(wire.data).map_err(serde::de::Error::custom)?,
156            ),
157            "FP64" => TensorData::F64(
158                serde_json::from_value(wire.data).map_err(serde::de::Error::custom)?,
159            ),
160            "INT32" => TensorData::I32(
161                serde_json::from_value(wire.data).map_err(serde::de::Error::custom)?,
162            ),
163            "INT64" => TensorData::I64(
164                serde_json::from_value(wire.data).map_err(serde::de::Error::custom)?,
165            ),
166            "UINT8" => {
167                TensorData::U8(serde_json::from_value(wire.data).map_err(serde::de::Error::custom)?)
168            }
169            "BOOL" => TensorData::Bool(
170                serde_json::from_value(wire.data).map_err(serde::de::Error::custom)?,
171            ),
172            "BYTES" => TensorData::Bytes(
173                serde_json::from_value(wire.data).map_err(serde::de::Error::custom)?,
174            ),
175            _ => serde_json::from_value(wire.data).map_err(serde::de::Error::custom)?,
176        };
177        Ok(Self {
178            dtype: wire.dtype,
179            shape: wire.shape,
180            data,
181        })
182    }
183}
184
185/// Tensor data variants supported by the core contract.
186#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
187#[serde(untagged)]
188pub enum TensorData {
189    /// 32-bit floating-point values.
190    F32(Vec<f32>),
191    /// 64-bit floating-point values.
192    F64(Vec<f64>),
193    /// 32-bit signed integer values.
194    I32(Vec<i32>),
195    /// 64-bit signed integer values.
196    I64(Vec<i64>),
197    /// 8-bit unsigned integer values.
198    U8(Vec<u8>),
199    /// Boolean values.
200    Bool(Vec<bool>),
201    /// Bytes values.
202    Bytes(Vec<Bytes>),
203}
204
205/// Adapter capability and executable authority declaration.
206#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
207pub struct InferenceDescriptor {
208    /// Stable adapter name.
209    pub name: String,
210    /// Human-readable adapter description.
211    pub description: String,
212    /// Serving protocol implemented by this adapter.
213    pub serving_protocol: ServingProtocol,
214    /// Executable permission envelope declared by the adapter.
215    pub envelope: rskit_tool::Envelope,
216}
217
218/// Model-serving runtime protocol.
219#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
220#[serde(rename_all = "snake_case")]
221#[non_exhaustive]
222pub enum ServingProtocol {
223    /// KServe v2 over HTTP.
224    KServeV2Http,
225    /// KServe v2 over gRPC.
226    KServeV2Grpc,
227    /// vLLM raw REST generation APIs.
228    VllmRest,
229    /// Hugging Face Text Generation Inference REST APIs.
230    TgiRest,
231    /// BentoML serving APIs.
232    BentoMl,
233    /// ONNX Runtime serving APIs.
234    OnnxRuntime,
235    /// TensorFlow Serving APIs.
236    TfServing,
237    /// Custom runtime protocol.
238    Custom,
239}
240
241/// Typed inference failure.
242#[derive(Debug, Error)]
243#[non_exhaustive]
244pub enum InferenceError {
245    /// Transport error from an injected client or adapter boundary.
246    #[error("transport: {0}")]
247    Transport(#[source] AppError),
248    /// Response decode or protocol mapping failure.
249    #[error("decode: {0}")]
250    Decode(String),
251    /// Serving runtime returned an unsuccessful response.
252    #[error("server: status={status}, body={body}")]
253    Server {
254        /// HTTP or protocol status code.
255        status: u16,
256        /// Runtime response body.
257        body: String,
258    },
259    /// Request was denied by an injected authorization decider.
260    #[error("authorization denied: {0}")]
261    Authorization(String),
262    /// Request or adapter selection input was invalid.
263    #[error("invalid input: {0}")]
264    InvalidInput(String),
265    /// Injected resilience policy failed the operation.
266    #[error("policy: {0}")]
267    Policy(String),
268    /// Requested adapter capability is not implemented yet.
269    #[error("not implemented: {0}")]
270    NotImplemented(&'static str),
271    /// Operation was cancelled.
272    #[error("cancelled")]
273    Cancelled,
274    /// Operation timed out.
275    #[error("timeout")]
276    Timeout,
277}
278
279impl From<AppError> for InferenceError {
280    fn from(value: AppError) -> Self {
281        match value.code() {
282            ErrorCode::Timeout => Self::Timeout,
283            ErrorCode::Cancelled => Self::Cancelled,
284            ErrorCode::ExternalService
285            | ErrorCode::ServiceUnavailable
286            | ErrorCode::ConnectionFailed => Self::Transport(value),
287            ErrorCode::InvalidInput | ErrorCode::MissingField | ErrorCode::InvalidFormat => {
288                Self::InvalidInput(value.to_string())
289            }
290            _ => Self::Policy(value.to_string()),
291        }
292    }
293}
294
295impl From<InferenceError> for AppError {
296    fn from(value: InferenceError) -> Self {
297        match value {
298            InferenceError::Timeout => AppError::timeout("inference"),
299            InferenceError::Cancelled => AppError::cancelled("inference"),
300            InferenceError::Authorization(reason) => AppError::forbidden(reason),
301            InferenceError::InvalidInput(reason) => AppError::new(ErrorCode::InvalidInput, reason),
302            InferenceError::Server { status, body } => AppError::new(
303                ErrorCode::ExternalService,
304                format!("inference runtime returned status {status}: {body}"),
305            ),
306            InferenceError::Transport(error) => AppError::new(
307                ErrorCode::ExternalService,
308                format!("inference transport failed: {}", error.message()),
309            )
310            .with_cause(error),
311            other => AppError::new(ErrorCode::ExternalService, other.to_string()),
312        }
313    }
314}
rskit_inference/types.rs

rskit_inference/
types.rs
