rill_runtime/handler/
builtin.rs1use serde::Deserialize;
9use serde_json::Value;
10
11use crate::package::LoadedModelPack;
12use crate::server::{InvokeError, InvokeErrorKind, InvokeHandler};
13
14pub const LINEAR_REGRESSION_CAPABILITY: &str = "rillml.linearRegression.predict";
15
16#[derive(Debug, Deserialize)]
17#[serde(rename_all = "camelCase", deny_unknown_fields)]
18struct LinearRegressionModel {
19 kind: String,
20 weights: Vec<f64>,
21 intercept: f64,
22}
23
24#[derive(Debug, Deserialize)]
25#[serde(deny_unknown_fields)]
26struct LinearRegressionInput {
27 features: Vec<f64>,
28}
29
30#[derive(Debug, Clone)]
32pub struct LinearRegressionInvokeHandler {
33 weights: Vec<f64>,
34 intercept: f64,
35}
36
37impl LinearRegressionInvokeHandler {
38 pub fn from_pack(pack: &LoadedModelPack) -> Result<Self, String> {
39 if pack.manifest.capabilities.as_slice() != [LINEAR_REGRESSION_CAPABILITY] {
40 return Err(format!(
41 "standalone runtime requires exactly the {LINEAR_REGRESSION_CAPABILITY} capability"
42 ));
43 }
44 let model: LinearRegressionModel = serde_json::from_value(pack.model.clone())
45 .map_err(|error| format!("invalid linear-regression model: {error}"))?;
46 if model.kind != "linearRegression" {
47 return Err("unsupported built-in model kind".into());
48 }
49 if model.weights.is_empty() || model.weights.len() > 65_536 {
50 return Err("linear-regression weights must contain 1..=65536 values".into());
51 }
52 if !model.intercept.is_finite() || model.weights.iter().any(|value| !value.is_finite()) {
53 return Err("linear-regression model values must be finite".into());
54 }
55 Ok(Self {
56 weights: model.weights,
57 intercept: model.intercept,
58 })
59 }
60}
61
62impl InvokeHandler for LinearRegressionInvokeHandler {
63 fn invoke(&self, capability: &str, input: &Value) -> Result<Value, InvokeError> {
64 if capability != LINEAR_REGRESSION_CAPABILITY {
65 return Err(InvokeError::with_detail(
69 InvokeErrorKind::Internal,
70 format!("built-in handler received unsupported capability: {capability}"),
71 ));
72 }
73 let input: LinearRegressionInput = serde_json::from_value(input.clone()).map_err(|e| {
74 InvokeError::with_detail(
75 InvokeErrorKind::Internal,
76 format!("invalid linear-regression input: {e}"),
77 )
78 })?;
79 if input.features.len() != self.weights.len() {
80 return Err(InvokeError::with_detail(
81 InvokeErrorKind::Internal,
82 format!(
83 "expected {} features, received {}",
84 self.weights.len(),
85 input.features.len()
86 ),
87 ));
88 }
89 if input.features.iter().any(|value| !value.is_finite()) {
90 return Err(InvokeError::with_detail(
91 InvokeErrorKind::Internal,
92 "linear-regression input values must be finite",
93 ));
94 }
95 let prediction = self
96 .weights
97 .iter()
98 .zip(&input.features)
99 .try_fold(self.intercept, |sum, (weight, feature)| {
100 let next = sum + weight * feature;
101 next.is_finite().then_some(next)
102 })
103 .ok_or_else(|| {
104 InvokeError::with_detail(
105 InvokeErrorKind::Internal,
106 "linear-regression prediction overflowed",
107 )
108 })?;
109 Ok(serde_json::json!({ "prediction": prediction }))
110 }
111}