cortexkit_role_classifier/
describe.rs

1//! `role.describe`: the module's limits and every model it serves.
2//!
3//! Resolved from the runner's catalog plus its own routing pins, so a
4//! caller validates a manifest at install time against what this runner
5//! will actually do, and can name the outside host its items are sent to.
6
7use std::collections::BTreeSet;
8
9use serde::{Deserialize, Serialize};
10use serde_json::Number;
11
12use crate::PROVIDES;
13
14/// The answer to `role.describe`: `{role, limits, models}`.
15#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
16#[non_exhaustive]
17pub struct RoleDescribe {
18    /// Always [`PROVIDES`], `classifier/v1`.
19    pub role: String,
20    pub limits: Limits,
21    pub models: Vec<ModelEntry>,
22}
23
24impl RoleDescribe {
25    pub fn new(limits: Limits, models: Vec<ModelEntry>) -> Self {
26        Self {
27            role: PROVIDES.to_owned(),
28            limits,
29            models,
30        }
31    }
32
33    /// The entry for `model`, if this runner serves it.
34    pub fn model(&self, model: &str) -> Option<&ModelEntry> {
35        self.models.iter().find(|entry| entry.model == model)
36    }
37}
38
39/// The module's own limits, across every model.
40#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize)]
41#[non_exhaustive]
42pub struct Limits {
43    /// Most items in one batch.
44    pub max_items: u64,
45    /// Largest `classify.run` request body, in bytes.
46    pub max_request_bytes: u64,
47}
48
49impl Limits {
50    pub fn new(max_items: u64, max_request_bytes: u64) -> Self {
51        Self {
52            max_items,
53            max_request_bytes,
54        }
55    }
56}
57
58/// One served model.
59#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
60#[non_exhaustive]
61pub struct ModelEntry {
62    /// The catalog id, `<provider>/<model>`.
63    pub model: String,
64    /// Most questions in one batch: the tighter of the catalog's and the
65    /// module's limits.
66    pub max_questions: u64,
67    /// Most items in one batch: the tighter of the catalog's and the
68    /// module's limits, so never above `limits.max_items`.
69    pub max_items: u64,
70    /// The model's context, in tokens, from the catalog. The provider
71    /// truncates `state` beyond it without saying so. Absent when the catalog
72    /// states no exact figure (a provider that says only "64k" gets no
73    /// guessed number); absent is not zero.
74    #[serde(default, skip_serializing_if = "Option::is_none")]
75    pub context_tokens: Option<u64>,
76    /// The model's image support; `null` when it takes no images. Always
77    /// encoded, never omitted, so "no images" is stated rather than implied.
78    pub images: Option<ImageSupport>,
79    /// The catalog price; absent when the model is unpriced.
80    #[serde(default, skip_serializing_if = "Option::is_none")]
81    pub price: Option<Price>,
82    /// The provider host item text is sent to, e.g. `api.typesafe.ai`.
83    pub egress_host: String,
84}
85
86impl ModelEntry {
87    /// `context_tokens` takes a number, or `None` when the catalog states
88    /// no exact figure.
89    pub fn new(
90        model: impl Into<String>,
91        max_questions: u64,
92        max_items: u64,
93        context_tokens: impl Into<Option<u64>>,
94        egress_host: impl Into<String>,
95    ) -> Self {
96        Self {
97            model: model.into(),
98            max_questions,
99            max_items,
100            context_tokens: context_tokens.into(),
101            images: None,
102            price: None,
103            egress_host: egress_host.into(),
104        }
105    }
106
107    pub fn with_images(mut self, images: ImageSupport) -> Self {
108        self.images = Some(images);
109        self
110    }
111
112    pub fn with_price(mut self, price: Price) -> Self {
113        self.price = Some(price);
114        self
115    }
116
117    /// Whether the model has a price for at least one token direction.
118    pub fn is_priced(&self) -> bool {
119        self.price.as_ref().is_some_and(|price| {
120            price.input_per_mtok_usd.is_some() || price.output_per_mtok_usd.is_some()
121        })
122    }
123}
124
125/// A model's image support: at most `max` images per item.
126#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize)]
127#[non_exhaustive]
128pub struct ImageSupport {
129    pub max: u64,
130}
131
132impl ImageSupport {
133    pub fn new(max: u64) -> Self {
134        Self { max }
135    }
136}
137
138/// A catalog price, in USD per million tokens. A direction the catalog does
139/// not price is absent, never zero.
140#[derive(Clone, Debug, Default, PartialEq, Deserialize, Serialize)]
141#[non_exhaustive]
142pub struct Price {
143    #[serde(default, skip_serializing_if = "Option::is_none")]
144    pub input_per_mtok_usd: Option<Number>,
145    #[serde(default, skip_serializing_if = "Option::is_none")]
146    pub output_per_mtok_usd: Option<Number>,
147}
148
149impl Price {
150    pub fn new() -> Self {
151        Self::default()
152    }
153
154    pub fn with_input_per_mtok_usd(mut self, usd: Number) -> Self {
155        self.input_per_mtok_usd = Some(usd);
156        self
157    }
158
159    pub fn with_output_per_mtok_usd(mut self, usd: Number) -> Self {
160        self.output_per_mtok_usd = Some(usd);
161        self
162    }
163}
164
165/// Why a `role.describe` answer is not usable.
166#[derive(Clone, Debug, PartialEq, Eq)]
167pub enum DescribeProblem {
168    /// `role` is not `classifier/v1`.
169    WrongRole(String),
170    /// The answer lists no model.
171    NoModels,
172    /// A model is listed twice.
173    DuplicateModel(String),
174    /// A model states a zero limit or a zero context, names no egress host,
175    /// or allows more items than the module does. Names the model and member.
176    BadModel { model: String, member: &'static str },
177    /// A module limit is zero.
178    BadLimits(&'static str),
179}
180
181impl std::fmt::Display for DescribeProblem {
182    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
183        match self {
184            Self::WrongRole(role) => write!(f, "role is {role}, not {PROVIDES}"),
185            Self::NoModels => f.write_str("models is empty"),
186            Self::DuplicateModel(model) => write!(f, "{model} is listed twice"),
187            Self::BadModel { model, member } => write!(f, "{model}: bad {member}"),
188            Self::BadLimits(member) => write!(f, "limits.{member} is zero"),
189        }
190    }
191}
192
193/// Check a `role.describe` answer before relying on it.
194pub fn check_describe(describe: &RoleDescribe) -> Result<(), DescribeProblem> {
195    if describe.role != PROVIDES {
196        return Err(DescribeProblem::WrongRole(describe.role.clone()));
197    }
198    if describe.limits.max_items == 0 {
199        return Err(DescribeProblem::BadLimits("max_items"));
200    }
201    if describe.limits.max_request_bytes == 0 {
202        return Err(DescribeProblem::BadLimits("max_request_bytes"));
203    }
204    if describe.models.is_empty() {
205        return Err(DescribeProblem::NoModels);
206    }
207    let mut seen = BTreeSet::new();
208    for entry in &describe.models {
209        if !seen.insert(entry.model.as_str()) {
210            return Err(DescribeProblem::DuplicateModel(entry.model.clone()));
211        }
212        let bad = |member| {
213            Err(DescribeProblem::BadModel {
214                model: entry.model.clone(),
215                member,
216            })
217        };
218        if entry.model.is_empty() {
219            return bad("model");
220        }
221        if entry.max_questions == 0 {
222            return bad("max_questions");
223        }
224        if entry.max_items == 0 || entry.max_items > describe.limits.max_items {
225            return bad("max_items");
226        }
227        // An absent context is allowed (the catalog states no exact
228        // figure); a stated zero is not.
229        if entry.context_tokens == Some(0) {
230            return bad("context_tokens");
231        }
232        if entry.images.is_some_and(|images| images.max == 0) {
233            return bad("images");
234        }
235        if entry.egress_host.is_empty() {
236            return bad("egress_host");
237        }
238    }
239    Ok(())
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245
246    fn describe() -> RoleDescribe {
247        RoleDescribe::new(
248            Limits::new(100, 1 << 20),
249            vec![
250                ModelEntry::new(
251                    "cloudflare-workers-ai/@cf/cloudflare/clef",
252                    64,
253                    100,
254                    8192,
255                    "api.cloudflare.com",
256                )
257                .with_images(ImageSupport::new(4)),
258                ModelEntry::new("typesafe/jev-1.13.0", 32, 50, 8192, "api.typesafe.ai"),
259            ],
260        )
261    }
262
263    #[test]
264    fn a_complete_answer_passes() {
265        assert_eq!(check_describe(&describe()), Ok(()));
266    }
267
268    #[test]
269    fn missing_images_encodes_as_null() {
270        let text = serde_json::to_string(&describe().models[1]).unwrap();
271        assert!(text.contains(r#""images":null"#), "{text}");
272        assert!(!text.contains("price"), "{text}");
273    }
274
275    #[test]
276    fn an_absent_context_is_omitted_and_passes() {
277        let mut answer = describe();
278        answer.models[1].context_tokens = None;
279        assert_eq!(check_describe(&answer), Ok(()));
280        let text = serde_json::to_string(&answer.models[1]).unwrap();
281        assert!(!text.contains("context_tokens"), "{text}");
282        let decoded: ModelEntry = serde_json::from_str(&text).unwrap();
283        assert_eq!(decoded.context_tokens, None);
284    }
285
286    #[test]
287    fn a_model_allowing_more_items_than_the_module_is_refused() {
288        let mut answer = describe();
289        answer.models[1].max_items = 101;
290        assert_eq!(
291            check_describe(&answer),
292            Err(DescribeProblem::BadModel {
293                model: "typesafe/jev-1.13.0".into(),
294                member: "max_items"
295            })
296        );
297        let mut answer = describe();
298        answer.models[0].context_tokens = Some(0);
299        assert!(matches!(
300            check_describe(&answer),
301            Err(DescribeProblem::BadModel {
302                member: "context_tokens",
303                ..
304            })
305        ));
306    }
307}
cortexkit_role_classifier/describe.rs

cortexkit_role_classifier/
describe.rs
