mecha10_core/

model.rs

//! Model configuration utilities for AI nodes
//!
//! This module provides shared utilities for loading and managing ONNX model configurations.
//! It's used by both the CLI (for model management) and AI nodes (for runtime inference).
//!
//! # Design
//!
//! Model configurations are stored in `models/catalog/<name>/config.json` and include:
//! - Model metadata (name, task, repo, filename)
//! - Preprocessing parameters (input_size, mean, std, channel_order)
//! - Class labels configuration
//!
//! # Usage
//!
//! ```no_run
//! use mecha10_core::model::{load_model_config, load_labels};
//!
//! // Load complete model configuration
//! let config = load_model_config("mobilenet-v2")?;
//!
//! // Load class labels
//! let labels = load_labels("mobilenet-v2")?;
//! # Ok::<(), anyhow::Error>(())
//! ```

use anyhow::{Context, Result};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

pub const CATALOG_MODELS_SUBDIR: &str = "models/catalog";

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelConfig {
    pub name: String,
    pub task: String,
    pub repo: String,
    pub filename: String,
    pub input_size: [u32; 2],
    pub preprocessing: PreprocessingConfig,
    pub num_classes: usize,
    pub labels_file: String,
    #[serde(default)]
    pub custom_labels: CustomLabelsConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PreprocessingConfig {
    pub mean: [f32; 3],
    pub std: [f32; 3],
    pub channel_order: String,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CustomLabelsConfig {
    #[serde(default)]
    pub enabled: bool,
    #[serde(default)]
    pub additional_classes: Vec<String>,
}

/// Where a node should resolve a model it references from.
///
/// Shared by every built-in node that points at an on-disk model this way (`object-detector`,
/// `image-classifier`, `rl` - see [`NodeModelConfig`]).
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum ModelSource {
    /// The public catalog of example models.
    #[default]
    Catalog,
    /// The private per-project model registry.
    Registry,
}

/// A node's reference to the on-disk model it should load - shared/reserved across every
/// built-in node that has this concept (`object-detector`, `image-classifier`, `rl`), moved
/// onto [`crate::node_config_base::NodeConfigBase`] (LAB-4384 follow-up) the same way
/// `relay_topics` already was, rather than each of those nodes re-declaring its own
/// near-identical `ModelConfig` struct.
///
/// This is deliberately a *different* type from [`ModelConfig`] above: that one is the model
/// catalog's own on-disk metadata (`models/catalog/<name>/config.json` - task, repo, filename,
/// preprocessing, labels); this one is a *node's pointer* to a model (which one, where its file
/// lives, where to resolve it from), a completely different concept that happens to share the
/// word "model".
///
/// `name`/`path` default to the empty string rather than this type baking in any one node's
/// real default value (object-detector defaults to `"yolov8n"`, rl to `"find-target"`,
/// image-classifier to `"mobilenet-v2"` - genuinely different per node). Each node resolves its
/// own real default for an empty value itself (see each node's own `model()` accessor on its
/// `Config` type) - this mirrors exactly how `#[serde(default)]` on this struct's own fields
/// already behaved before this type was shared (calling `Default::default()`, not a per-node
/// default function, when the field - or the whole `model` key - is merely absent).
///
/// `confidence_threshold` only has meaning for `image-classifier` today (object-detector/rl
/// ignore it) - it lives here rather than as a `image-classifier`-only extra field so this
/// struct stays the single shared shape for the "which model does this node use" concept across
/// all three nodes rather than splintering into near-identical siblings again.
///
/// `provider` only has meaning for `llm-command` (object-detector/image-classifier/rl ignore
/// it) - it used to be a separate field on `NodeConfigBase` alongside `model`, but "which
/// backend resolves this model" is part of the same "which model does this node use" concept,
/// so it consolidated into this one struct instead of two separate `Option` fields on the base.
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct NodeModelConfig {
    /// Name of the model to load. Empty means "not set" - the owning node resolves its own
    /// default in that case.
    #[serde(default)]
    pub name: String,

    /// Which LLM backend to query (`llm-command`-specific; ignored by other nodes using this
    /// type). `None` means "not set" - the owning node resolves its own default in that case.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub provider: Option<Provider>,

    /// Filesystem path to the model file. Empty means "not set" - the owning node resolves its
    /// own default in that case.
    #[serde(default)]
    pub path: String,

    /// Where to resolve the model from.
    #[serde(default)]
    pub source: ModelSource,

    /// Specific model version to use (registry-sourced models only); latest if omitted.
    #[serde(default)]
    pub version: Option<i32>,

    /// Minimum confidence score a classification result must meet to be reported
    /// (`image-classifier`-specific; ignored by other nodes using this type).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub confidence_threshold: Option<f32>,
}

/// Which LLM backend a node should query - shared/reserved the same way [`NodeModelConfig`] is,
/// for `llm-command`'s analogous "which model does this node use" concern (LAB-4384 follow-up).
/// `llm-command` is the only node with this concept today, so unlike [`NodeModelConfig`] there's
/// no cross-node shape to reconcile - this just lives in `mecha10-core` (rather than
/// `mecha10-nodes-llm-command`) so [`crate::node_config_base::NodeConfigBase`] can reference it.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default, JsonSchema)]
#[serde(rename_all = "lowercase")]
pub enum Provider {
    /// OpenAI's hosted API.
    Openai,
    /// Anthropic's hosted API.
    Claude,
    /// A local Ollama server.
    #[default]
    Local,
}

crate::enum_display_via_str!(Provider {
    Openai => "openai",
    Claude => "claude",
    Local => "local",
});

/// Load model configuration from disk
///
/// Loads the complete model configuration from `models/catalog/<name>/config.json`.
///
/// # Arguments
///
/// * `model_name` - Name of the model (e.g., "mobilenet-v2")
///
/// # Returns
///
/// The complete `ModelConfig` with all metadata and preprocessing parameters.
///
/// # Errors
///
/// Returns an error if:
/// - The config file doesn't exist
/// - The config file is not valid JSON
/// - The config file is missing required fields
///
/// # Example
///
/// ```no_run
/// use mecha10_core::model::load_model_config;
///
/// let config = load_model_config("mobilenet-v2")?;
/// println!("Model: {} ({})", config.name, config.task);
/// println!("Input size: {:?}", config.input_size);
/// println!("Preprocessing: mean={:?}, std={:?}", config.preprocessing.mean, config.preprocessing.std);
/// # Ok::<(), anyhow::Error>(())
/// ```
pub fn load_model_config(model_name: &str) -> Result<ModelConfig> {
    let config_path = PathBuf::from(CATALOG_MODELS_SUBDIR)
        .join(model_name)
        .join("config.json");

    if !config_path.exists() {
        anyhow::bail!(
            "Model config not found for '{}' at {}. Run: mecha10 model pull {}",
            model_name,
            config_path.display(),
            model_name
        );
    }

    let content = std::fs::read_to_string(&config_path)
        .context(format!("Failed to read model config: {}", config_path.display()))?;

    let config: ModelConfig = serde_json::from_str(&content)
        .context(format!("Failed to parse model config.json: {}", config_path.display()))?;

    Ok(config)
}

/// Load class labels from model directory
///
/// Loads class labels from `models/catalog/<name>/labels.txt` with support for custom labels.
/// If custom labels are enabled in the model config, they will be appended to the base labels.
///
/// # Arguments
///
/// * `model_name` - Name of the model (e.g., "mobilenet-v2")
///
/// # Returns
///
/// A vector of class label strings. The index in the vector corresponds to the class ID.
///
/// # Errors
///
/// Returns an error if:
/// - The labels file doesn't exist
/// - The labels file cannot be read
/// - The model config cannot be loaded (when checking for custom labels)
///
/// # Example
///
/// ```no_run
/// use mecha10_core::model::load_labels;
///
/// let labels = load_labels("mobilenet-v2")?;
/// println!("Loaded {} class labels", labels.len());
/// println!("First label: {}", labels[0]);
/// # Ok::<(), anyhow::Error>(())
/// ```
pub fn load_labels(model_name: &str) -> Result<Vec<String>> {
    let labels_path = PathBuf::from(CATALOG_MODELS_SUBDIR).join(model_name).join("labels.txt");

    if !labels_path.exists() {
        anyhow::bail!(
            "Labels file not found for model '{}' at {}. Run: mecha10 model pull {}",
            model_name,
            labels_path.display(),
            model_name
        );
    }

    tracing::info!("📋 Loading labels from: {}", labels_path.display());

    let content = std::fs::read_to_string(&labels_path)
        .context(format!("Failed to read labels file: {}", labels_path.display()))?;

    let mut labels: Vec<String> = content.lines().map(|s| s.trim().to_string()).collect();

    let config_path = PathBuf::from(CATALOG_MODELS_SUBDIR)
        .join(model_name)
        .join("config.json");
    if config_path.exists() {
        if let Ok(config_content) = std::fs::read_to_string(&config_path) {
            if let Ok(config) = serde_json::from_str::<ModelConfig>(&config_content) {
                if config.custom_labels.enabled {
                    let num_custom = config.custom_labels.additional_classes.len();
                    labels.extend(config.custom_labels.additional_classes);
                    tracing::info!("✅ Added {} custom labels (total: {} labels)", num_custom, labels.len());
                }
            }
        }
    }

    tracing::info!("✅ Loaded {} labels", labels.len());
    Ok(labels)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_model_config_serialization() {
        let config = ModelConfig {
            name: "test-model".to_string(),
            task: "image-classification".to_string(),
            repo: "test/repo".to_string(),
            filename: "model.onnx".to_string(),
            input_size: [224, 224],
            preprocessing: PreprocessingConfig {
                mean: [0.485, 0.456, 0.406],
                std: [0.229, 0.224, 0.225],
                channel_order: "RGB".to_string(),
            },
            num_classes: 1000,
            labels_file: "labels.txt".to_string(),
            custom_labels: CustomLabelsConfig::default(),
        };

        let json = serde_json::to_string(&config).unwrap();
        let deserialized: ModelConfig = serde_json::from_str(&json).unwrap();

        assert_eq!(config.name, deserialized.name);
        assert_eq!(config.input_size, deserialized.input_size);
        assert_eq!(config.preprocessing.mean, deserialized.preprocessing.mean);
    }

    /// Regression test for `LAB-2073`: a fresh `mecha10 model pull
    /// <name>`-style write must produce files this module's own path
    /// construction can find. `mecha10_models::ModelService::catalog_model_dir`
    /// (see `packages/mecha10-models/src/services/model_service/mod.rs`)
    /// writes catalog pulls to `<models_dir>/catalog/<name>/` (`LAB-1802`);
    /// this asserts `load_model_config`/`load_labels`'s path construction
    /// matches that layout exactly rather than drifting back to the
    /// pre-`LAB-1802` flat `models/<name>/` path. `mecha10-models` can't be
    /// imported directly here (see `CATALOG_MODELS_SUBDIR`'s doc comment), so
    /// this pins the literal `"catalog"` segment instead - mirroring
    /// `object-detector`'s
    /// `default_model_path_matches_model_services_catalog_write_location`.
    #[test]
    fn catalog_paths_match_model_services_catalog_write_location() {
        assert_eq!(CATALOG_MODELS_SUBDIR, "models/catalog");

        let model_name = "mobilenet-v2";
        let config_path = PathBuf::from(CATALOG_MODELS_SUBDIR)
            .join(model_name)
            .join("config.json");
        let labels_path = PathBuf::from(CATALOG_MODELS_SUBDIR).join(model_name).join("labels.txt");

        assert_eq!(config_path, PathBuf::from("models/catalog/mobilenet-v2/config.json"));
        assert_eq!(labels_path, PathBuf::from("models/catalog/mobilenet-v2/labels.txt"));
    }
}