mecha10-models 0.9.0

//! Registry model artifact download (`LAB-1802`)
//!
//! Distinct from `download.rs`'s public-catalog HuggingFace pulls: this hits the
//! private, versioned, per-project model registry
//! (`packages/mecha10-api/src/models/registry`) via the same by-name -> version ->
//! download resolution the `mecha10 model set` CLI command
//! (`packages/cli/src/handlers/model.rs`) already used ad hoc. Pulling that
//! resolution into `ModelService` lets any caller that knows a robot/environment
//! is source-aware (`node-runner`'s preflight, the CLI's bundled-node runner) fetch
//! a *specific* registry version instead of only being able to assume every model
//! name is a catalog entry.
//!
//! Artifacts are cached at `models_dir/registry/<name>/<version>/model.onnx` -
//! namespaced separately from the public catalog's `models_dir/catalog/<name>/`
//! tree so a name collision between the two systems can never silently overwrite
//! the wrong artifact. See the module docs on [`super::ModelService`].

use super::ModelService;
use std::path::PathBuf;

impl ModelService {
    /// `models_dir/registry/<model_name>/<version>/` - a registry model's
    /// namespaced, per-version subdirectory (`LAB-1802`). Pure path computation
    /// (no I/O), so it's available regardless of the `model-download` feature -
    /// callers that only need to check/derive paths (e.g. `node-runner`'s
    /// preflight existence check) don't need the download machinery linked in.
    pub fn registry_model_dir(&self, model_name: &str, version: i32) -> PathBuf {
        self.models_dir
            .join("registry")
            .join(model_name)
            .join(version.to_string())
    }
}

#[cfg(feature = "model-download")]
mod download {
    use super::ModelService;
    use anyhow::{Context, Result};
    use indicatif::{ProgressBar, ProgressStyle};
    use serde::Deserialize;
    use std::path::PathBuf;
    use tokio::fs;

    #[derive(Debug, Deserialize)]
    struct ModelLookupResponse {
        id: String,
    }

    impl ModelService {
        /// Resolve `model_name`'s registry `Model` id via `mecha10-api`, then
        /// download `version`'s artifact and cache it at
        /// `registry_model_dir(model_name, version)/model.onnx`. Returns the
        /// cached path immediately (without a network round-trip) if it's
        /// already present.
        pub async fn pull_registry_version(
            &self,
            base_url: &str,
            api_key: &str,
            model_name: &str,
            version: i32,
            progress: Option<&ProgressBar>,
        ) -> Result<PathBuf> {
            let model_dir = self.registry_model_dir(model_name, version);
            fs::create_dir_all(&model_dir).await?;
            let output_path = model_dir.join("model.onnx");

            if output_path.exists() {
                if let Some(pb) = progress {
                    pb.set_message(format!("Model '{model_name}' v{version} already cached"));
                }
                return Ok(output_path);
            }

            let base = base_url.trim_end_matches('/');
            let client = reqwest::Client::builder()
                .user_agent("mecha10-cli")
                .build()
                .context("Failed to build HTTP client")?;

            let lookup: ModelLookupResponse = client
                .get(format!("{base}/api/models/by-name/{model_name}"))
                .bearer_auth(api_key)
                .send()
                .await
                .context(format!("failed to reach the model registry for '{model_name}'"))?
                .error_for_status()
                .context(format!("model '{model_name}' was not found in the registry"))?
                .json()
                .await
                .context("model registry returned an unexpected response")?;

            if let Some(pb) = progress {
                pb.set_style(
                    ProgressStyle::default_spinner()
                        .template("{spinner:.green} {msg}")
                        .unwrap(),
                );
                pb.set_message(format!("Downloading {model_name} v{version} from registry"));
            }

            let download_url = format!("{base}/api/models/{}/versions/number/{version}/download", lookup.id);

            tracing::info!("Downloading registry model from mecha10-api: {}", download_url);

            let response = client
                .get(&download_url)
                .bearer_auth(api_key)
                .send()
                .await
                .context(format!("failed to download registry model '{model_name}' v{version}"))?
                .error_for_status()
                .context(format!("registry rejected the download for '{model_name}' v{version}"))?;

            let mut file = tokio::fs::File::create(&output_path).await?;
            let mut stream = response.bytes_stream();

            use futures_util::StreamExt;
            use tokio::io::AsyncWriteExt;

            while let Some(chunk) = stream.next().await {
                let chunk = chunk.context("Error downloading registry model chunk")?;
                file.write_all(&chunk).await?;
            }
            file.flush().await?;

            if let Some(pb) = progress {
                pb.set_message(format!("Downloaded {model_name} v{version} successfully"));
            }

            Ok(output_path)
        }
    }
}

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

    fn service() -> ModelService {
        ModelService::with_models_dir(PathBuf::from("models")).unwrap()
    }

    #[test]
    fn registry_model_dir_is_namespaced_separately_from_catalog() {
        let service = service();
        let registry_dir = service.registry_model_dir("rover-detector", 3);
        let catalog_dir = service.catalog_model_dir("rover-detector");

        assert_eq!(registry_dir, PathBuf::from("models/registry/rover-detector/3"));
        assert_eq!(catalog_dir, PathBuf::from("models/catalog/rover-detector"));
        assert_ne!(registry_dir, catalog_dir);
    }

    #[test]
    fn registry_model_dir_is_namespaced_by_version() {
        let service = service();
        let v1 = service.registry_model_dir("rover-detector", 1);
        let v2 = service.registry_model_dir("rover-detector", 2);

        assert_ne!(v1, v2);
    }
}