cortexkit-role-classifier 0.4.2

//! `classify.run` requests and their validators.
//!
//! A runner validates a request completely before it records anything or
//! calls a provider: [`parse_request`] decodes the request text naming the
//! first malformed field, [`check_request_bytes`] checks its size,
//! [`check_catalog_row`] the model's catalog row, and [`check_request`]
//! every limit the model's `role.describe` entry states. Each refusal
//! carries its code and `detail.field`.

use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use serde_json::{Number, Value};

use cortexkit_model_batch::{
    admission::{
        check_batch_id, check_catalog_row as check_row, check_max_cost_usd, check_request_size,
    },
    image::{check_images, ImageLimits},
};

pub use cortexkit_model_batch::admission::{
    effective_ceiling, is_batch_id, CatalogRow, BATCH_ID_MAX_LEN,
};

use crate::{
    describe::{Limits, ModelEntry},
    errors::Refusal,
    question::{bounded, check_question, Image, Question},
};

/// The catalog `kind` of a classifier model.
pub const CLASSIFIER_KIND: &str = "classifier";

/// The `classify.run` request.
///
/// Non-exhaustive so later optional members are additive: use
/// [`ClassifyRequest::new`] and the `with_*` setters, or decode one.
#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
#[non_exhaustive]
pub struct ClassifyRequest {
    /// The idempotency key: 1 to [`BATCH_ID_MAX_LEN`] printable ASCII
    /// characters.
    pub batch_id: String,
    /// The catalog id, `<provider>/<model>`. Never substituted.
    pub model: String,
    /// The questions, by caller-chosen id, in the order the caller wrote
    /// them.
    pub questions: IndexMap<String, Question>,
    /// The items, answered in this order.
    pub items: Vec<Item>,
    /// The batch's spend ceiling. Not part of the body identity; a lower
    /// value on a re-send becomes the batch's recorded ceiling (see
    /// [`effective_ceiling`]).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_cost_usd: Option<Number>,
}

impl ClassifyRequest {
    pub fn new(
        batch_id: impl Into<String>,
        model: impl Into<String>,
        questions: IndexMap<String, Question>,
        items: Vec<Item>,
    ) -> Self {
        Self {
            batch_id: batch_id.into(),
            model: model.into(),
            questions,
            items,
            max_cost_usd: None,
        }
    }

    pub fn with_max_cost_usd(mut self, max_cost_usd: Number) -> Self {
        self.max_cost_usd = Some(max_cost_usd);
        self
    }
}

/// One item: `{state, images?}`.
#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
#[non_exhaustive]
pub struct Item {
    /// A string, an object or an array, forwarded to the provider verbatim.
    /// Null, numbers and booleans are refused.
    pub state: Value,
    /// Only where the model's catalog row allows images.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub images: Option<Vec<Image>>,
}

impl Item {
    pub fn new(state: impl Into<Value>) -> Self {
        Self {
            state: state.into(),
            images: None,
        }
    }

    pub fn with_images(mut self, images: Vec<Image>) -> Self {
        self.images = Some(images);
        self
    }
}

/// Whether `state` is a kind the role forwards: a string, an object or an
/// array.
pub fn is_state(state: &Value) -> bool {
    matches!(state, Value::String(_) | Value::Object(_) | Value::Array(_))
}

/// Decode a `classify.run` request from its JSON text, refusing a malformed
/// one `invalid_params` naming the first malformed field. Decoding from the
/// text keeps the order of the questions and of each choice's options.
///
/// The refusal names `params` only when no narrower path applies: the text
/// is not JSON, or not an object. A problem deep inside an otherwise
/// well-formed request names its own path (`items[2].images[0]`, say),
/// never `params`.
pub fn parse_request(text: &str) -> Result<ClassifyRequest, Refusal> {
    let value: Value = serde_json::from_str(text)
        .map_err(|e| Refusal::invalid_params("params", format!("not JSON: {e}")))?;
    let Value::Object(params) = &value else {
        return Err(Refusal::invalid_params("params", "params is an object"));
    };
    let malformed = |field: String, what: &str| Refusal::invalid_params(field, what.to_owned());
    for (field, what) in [
        ("batch_id", "a string"),
        ("model", "a string"),
        ("questions", "an object"),
        ("items", "an array"),
    ] {
        let ok = matches!(
            (params.get(field), what),
            (Some(Value::String(_)), "a string")
                | (Some(Value::Object(_)), "an object")
                | (Some(Value::Array(_)), "an array")
        );
        if !ok {
            return Err(malformed(field.to_owned(), &format!("{field} is {what}")));
        }
    }
    match params.get("max_cost_usd") {
        None | Some(Value::Null) | Some(Value::Number(_)) => {}
        Some(_) => return Err(malformed("max_cost_usd".into(), "max_cost_usd is a number")),
    }
    for (id, question) in params["questions"].as_object().into_iter().flatten() {
        let Value::Object(question) = question else {
            return Err(malformed(
                format!("questions.{id}"),
                "a question is an object",
            ));
        };
        if !question.get("type").is_some_and(Value::is_string) {
            return Err(malformed(
                format!("questions.{id}.type"),
                "type is a string",
            ));
        }
        if !question.contains_key("instructions") {
            return Err(malformed(
                format!("questions.{id}.instructions"),
                "instructions is required",
            ));
        }
        if !matches!(
            question.get("criteria"),
            None | Some(Value::Null) | Some(Value::Array(_)) | Some(Value::Object(_))
        ) {
            return Err(malformed(
                format!("questions.{id}.criteria"),
                "criteria is an array or an object",
            ));
        }
    }
    for (i, item) in params["items"].as_array().into_iter().flatten().enumerate() {
        let Value::Object(item) = item else {
            return Err(malformed(format!("items[{i}]"), "an item is an object"));
        };
        if !item.get("state").is_some_and(is_state) {
            return Err(malformed(
                format!("items[{i}].state"),
                "state is a string, an object or an array",
            ));
        }
        match item.get("images") {
            None | Some(Value::Null) => {}
            Some(Value::Array(images)) => {
                for (j, image) in images.iter().enumerate() {
                    if serde_json::from_value::<Image>(image.clone()).is_err() {
                        return Err(malformed(
                            format!("items[{i}].images[{j}]"),
                            "an image is a data: URL or {content_type, base64}",
                        ));
                    }
                }
            }
            Some(_) => {
                return Err(malformed(
                    format!("items[{i}].images"),
                    "images is an array",
                ))
            }
        }
    }
    // Every member was checked above, so this decode does not fail for any
    // reason a narrower path could name; `params` covers what is left.
    serde_json::from_str(text).map_err(|e| Refusal::invalid_params("params", e.to_string()))
}

/// Check the request body's size against the module's limit. The refusal
/// names `params` and the limit: the size belongs to the request as a
/// whole, so no narrower path applies.
pub fn check_request_bytes(len: usize, limits: &Limits) -> Result<(), Refusal> {
    check_request_size(len, limits.max_request_bytes)
}

/// Check the model's catalog row: [`CatalogRow::Absent`] is `model_unknown`;
/// a row with another kind or no kind is `model_not_classifier`. A kindless
/// refusal names the model without inventing a kind.
pub fn check_catalog_row(model: &str, row: CatalogRow<'_>) -> Result<(), Refusal> {
    check_row(model, row, CLASSIFIER_KIND, |model, kind| match kind {
        Some(kind) => Refusal::model_not_classifier(model, kind),
        None => {
            let mut refusal = Refusal::model_not_classifier(model, "");
            refusal.message = format!("{model} is not a classifier");
            refusal.detail.kind = None;
            refusal
        }
    })
}

/// Check every limit of `request` against the model's `role.describe` entry:
/// the batch id, question count and each question, item count, each item's
/// state and images, and `max_cost_usd`. Returns the first refusal.
pub fn check_request(request: &ClassifyRequest, model: &ModelEntry) -> Result<(), Refusal> {
    check_batch_id(&request.batch_id)?;
    bounded(
        request.questions.len(),
        1,
        model.max_questions as usize,
        "questions",
        "questions",
    )?;
    for (id, question) in &request.questions {
        check_question(id, question)?;
    }
    bounded(
        request.items.len(),
        1,
        model.max_items as usize,
        "items",
        "items",
    )?;
    for (i, item) in request.items.iter().enumerate() {
        if !is_state(&item.state) {
            return Err(Refusal::invalid_params(
                format!("items[{i}].state"),
                "state is a string, an object or an array",
            ));
        }
        if let Some(images) = &item.images {
            let limits = model
                .images
                .map(|support| ImageLimits::workers_ai(support.max));
            check_images(
                &model.model,
                &format!("items[{i}].images"),
                images,
                limits.as_ref(),
            )?;
        }
    }
    check_max_cost_usd(request.max_cost_usd.as_ref())?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;
    use crate::{
        describe::ImageSupport,
        errors::{INVALID_PARAMS, MODEL_NOT_CLASSIFIER, MODEL_UNKNOWN},
        question::InlineImage,
    };

    fn model() -> ModelEntry {
        ModelEntry::new(
            "cloudflare-workers-ai/@cf/cloudflare/clef",
            2,
            3,
            8192,
            "api.cloudflare.com",
        )
        .with_images(ImageSupport::new(4))
    }

    fn request(items: Vec<Item>) -> ClassifyRequest {
        ClassifyRequest::new(
            "b-1",
            "cloudflare-workers-ai/@cf/cloudflare/clef",
            [("urgent".to_owned(), Question::noul("Is it urgent?"))].into(),
            items,
        )
    }

    fn field(result: Result<(), Refusal>) -> String {
        let refusal = result.unwrap_err();
        assert_eq!(refusal.code, INVALID_PARAMS);
        refusal.detail.field.unwrap()
    }

    #[test]
    fn a_state_of_null_number_or_boolean_is_refused_naming_its_index() {
        for bad in [json!(null), json!(3), json!(true)] {
            let items = vec![Item::new("ok"), Item::new(json!({"a": 1})), Item::new(bad)];
            assert_eq!(
                field(check_request(&request(items), &model())),
                "items[2].state"
            );
        }
        let items = vec![
            Item::new("s"),
            Item::new(json!({"a": 1})),
            Item::new(json!([1])),
        ];
        assert!(check_request(&request(items), &model()).is_ok());
        let refusal = parse_request(
            r#"{"batch_id":"b","model":"m","questions":{},"items":[{"state":"s"},{"state":null}]}"#,
        )
        .unwrap_err();
        assert_eq!(refusal.field(), Some("items[1].state"));
    }

    #[test]
    fn counts_are_refused_with_their_limit() {
        let items = (0..4).map(|i| Item::new(format!("s{i}"))).collect();
        let refusal = check_request(&request(items), &model()).unwrap_err();
        assert_eq!(refusal.field(), Some("items"));
        assert_eq!(refusal.detail.limit, Some(3));
        let mut three = request(vec![Item::new("s")]);
        for id in ["a", "b"] {
            three.questions.insert(id.into(), Question::noul("i"));
        }
        let refusal = check_request(&three, &model()).unwrap_err();
        assert_eq!(refusal.field(), Some("questions"));
        assert_eq!(refusal.detail.limit, Some(2));
        assert_eq!(field(check_request(&request(vec![]), &model())), "items");
    }

    #[test]
    fn images_are_refused_where_the_model_takes_none_or_too_many() {
        let image = Image::Inline(InlineImage::new("image/png", "AAAA"));
        let item = Item::new("s").with_images(vec![image.clone(); 5]);
        let refusal = check_request(&request(vec![item]), &model()).unwrap_err();
        assert_eq!(refusal.field(), Some("items[0].images"));
        assert_eq!(refusal.detail.limit, Some(4));
        let text_only = ModelEntry::new("typesafe/jev-1.13.0", 2, 3, 8192, "api.typesafe.ai");
        let item = Item::new("s").with_images(vec![image]);
        let refusal = check_request(&request(vec![item]), &text_only).unwrap_err();
        assert_eq!(refusal.field(), Some("items[0].images"));
        assert_eq!(refusal.detail.limit, Some(0));
        let item = Item::new("s").with_images(vec![Image::DataUrl("https://x/y.png".into())]);
        assert_eq!(
            field(check_request(&request(vec![item]), &model())),
            "items[0].images[0]"
        );
        // Malformed base64 is refused before any provider call, naming the
        // one image.
        let item = Item::new("s").with_images(vec![Image::Inline(InlineImage::new(
            "image/png",
            "not base64!",
        ))]);
        assert_eq!(
            field(check_request(&request(vec![item]), &model())),
            "items[0].images[0]"
        );
    }

    #[test]
    fn batch_id_and_max_cost_are_checked() {
        let mut bad = request(vec![Item::new("s")]);
        bad.batch_id = "tab\there".into();
        assert_eq!(field(check_request(&bad, &model())), "batch_id");
        bad.batch_id = "x".repeat(129);
        assert_eq!(field(check_request(&bad, &model())), "batch_id");
        let negative =
            request(vec![Item::new("s")]).with_max_cost_usd(Number::from_f64(-1.0).unwrap());
        assert_eq!(field(check_request(&negative, &model())), "max_cost_usd");
    }

    #[test]
    fn parse_names_the_malformed_field_and_keeps_question_order() {
        let refusal = parse_request(r#"{"model":"m","questions":{},"items":[]}"#).unwrap_err();
        assert_eq!(refusal.field(), Some("batch_id"));
        let refusal = parse_request(
            r#"{"batch_id":"b","model":"m","questions":{"q":{"instructions":"i"}},"items":[]}"#,
        )
        .unwrap_err();
        assert_eq!(refusal.field(), Some("questions.q.type"));
        let parsed = parse_request(
            r#"{"batch_id":"b","model":"m","questions":{"z":{"type":"noul","instructions":"i"},"a":{"type":"noul","instructions":"i"}},"items":[{"state":"s"}]}"#,
        )
        .unwrap();
        assert_eq!(parsed.questions.keys().collect::<Vec<_>>(), ["z", "a"]);
    }

    #[test]
    fn catalog_rows_are_pinned() {
        assert!(check_catalog_row("m", CatalogRow::Kind("classifier")).is_ok());
        assert_eq!(
            check_catalog_row("m", CatalogRow::Absent).unwrap_err().code,
            MODEL_UNKNOWN
        );
        let kindless = check_catalog_row("m", CatalogRow::Kindless).unwrap_err();
        assert_eq!(kindless.code, MODEL_NOT_CLASSIFIER);
        assert_eq!(kindless.detail.model.as_deref(), Some("m"));
        assert_eq!(kindless.detail.kind, None);
        let other_kind = check_catalog_row("m", CatalogRow::Kind("chat")).unwrap_err();
        assert_eq!(other_kind.code, MODEL_NOT_CLASSIFIER);
        assert_eq!(other_kind.detail.model.as_deref(), Some("m"));
        assert_eq!(other_kind.detail.kind.as_deref(), Some("chat"));
    }

    #[test]
    fn request_bytes_are_bounded() {
        let limits = Limits::new(10, 100);
        assert!(check_request_bytes(100, &limits).is_ok());
        let refusal = check_request_bytes(101, &limits).unwrap_err();
        assert_eq!(refusal.field(), Some("params"));
        assert_eq!(refusal.detail.limit, Some(100));
    }

    #[test]
    fn the_ceiling_tightens_and_never_loosens() {
        assert_eq!(effective_ceiling(Some(1.0), Some(5.0)), Some(1.0));
        assert_eq!(effective_ceiling(Some(1.0), Some(0.5)), Some(0.5));
        assert_eq!(effective_ceiling(None, Some(0.5)), Some(0.5));
        assert_eq!(effective_ceiling(Some(1.0), None), Some(1.0));
        assert_eq!(effective_ceiling(None, None), None);
    }

    #[test]
    fn a_tightened_ceiling_is_recorded_across_resends() {
        // High, then lower, then the original high again: each call records
        // its effective ceiling, so the third works under the lower value.
        let mut recorded = None;
        let mut worked_under = Vec::new();
        for sent in [5.0, 1.0, 5.0] {
            recorded = effective_ceiling(recorded, Some(sent));
            worked_under.push(recorded);
        }
        assert_eq!(worked_under, [Some(5.0), Some(1.0), Some(1.0)]);
    }

    #[test]
    fn params_is_named_only_when_no_narrower_path_applies() {
        for whole in ["", "not json", "[]", "null", "7"] {
            assert_eq!(
                parse_request(whole).unwrap_err().field(),
                Some("params"),
                "{whole:?}"
            );
        }
        // An empty object misses its first required member, which is a
        // narrower path than the request as a whole.
        assert_eq!(parse_request("{}").unwrap_err().field(), Some("batch_id"));
        // A malformed image deep in the request keeps the whole text from
        // decoding, and the refusal still names the one image.
        let refusal = parse_request(
            r#"{"batch_id":"b","model":"m","questions":{"q":{"type":"noul","instructions":"i"}},"items":[{"state":"s"},{"state":"t","images":["data:image/png;base64,AAAA",5]}]}"#,
        )
        .unwrap_err();
        assert_eq!(refusal.field(), Some("items[1].images[1]"));
        let refusal = parse_request(
            r#"{"batch_id":"b","model":"m","questions":{"q":{"type":"noul","instructions":"i","criteria":"yes"}},"items":[{"state":"s"}]}"#,
        )
        .unwrap_err();
        assert_eq!(refusal.field(), Some("questions.q.criteria"));
    }

    #[test]
    fn an_unknown_question_type_in_a_request_is_refused_naming_its_key() {
        let mut request = request(vec![Item::new("s")]);
        request.questions.insert(
            "priority".into(),
            Question::new(crate::question::QuestionType::parse("rank"), "How urgent?"),
        );
        assert_eq!(
            field(check_request(&request, &model())),
            "questions.priority.type"
        );
    }
}