From bc02e5ff31154e307816a7f49e748f1ff81d447b Mon Sep 17 00:00:00 2001 From: jack-edmonds-dd Date: Thu, 18 Jun 2026 13:34:38 -0400 Subject: [PATCH 1/2] feat(model-lab): add model-lab command group (SDK 0.32.0) Implements the Model Lab API surface added in SDK 0.32.0: Projects: - list: list projects with optional filter/tags/sort/pagination - get: fetch a project by ID - star / unstar: mark projects as starred - artifacts: list project artifacts - facet-keys / facet-values: explore project facets Runs: - list: list runs with rich filtering (project, status, tags, params, parent) - get: fetch a run by ID - delete: delete a run - pin / unpin: pin runs for easy access - artifacts: list run artifacts with optional path prefix filter - facet-keys / facet-values: explore run facets All 16 operations registered in UNSTABLE_OPS. 18 unit tests, all passing. Co-Authored-By: Claude Sonnet 4.6 --- src/client.rs | 19 +- src/commands/mod.rs | 1 + src/commands/model_lab.rs | 520 ++++++++++++++++++++++++++++++++++++++ src/config.rs | 1 + src/main.rs | 246 ++++++++++++++++++ 5 files changed, 786 insertions(+), 1 deletion(-) create mode 100644 src/commands/model_lab.rs diff --git a/src/client.rs b/src/client.rs index 95098fed..77d03cf0 100644 --- a/src/client.rs +++ b/src/client.rs @@ -398,6 +398,23 @@ static UNSTABLE_OPS: &[&str] = &[ "v2.trigger_investigation", // Cloud Cost Management — Anomalies (1) "v2.list_cost_anomalies", + // Model Lab (16) + "v2.delete_model_lab_run", + "v2.get_model_lab_artifact_content", + "v2.get_model_lab_project", + "v2.get_model_lab_run", + "v2.list_model_lab_project_artifacts", + "v2.list_model_lab_project_facet_keys", + "v2.list_model_lab_project_facet_values", + "v2.list_model_lab_projects", + "v2.list_model_lab_run_artifacts", + "v2.list_model_lab_run_facet_keys", + "v2.list_model_lab_run_facet_values", + "v2.list_model_lab_runs", + "v2.pin_model_lab_run", + "v2.star_model_lab_project", + "v2.unpin_model_lab_run", + "v2.unstar_model_lab_project", ]; // --------------------------------------------------------------------------- @@ -1303,7 +1320,7 @@ mod tests { #[test] fn test_unstable_ops_count() { - assert_eq!(UNSTABLE_OPS.len(), 171); + assert_eq!(UNSTABLE_OPS.len(), 187); } #[test] diff --git a/src/commands/mod.rs b/src/commands/mod.rs index a8bb41c9..a28ce14c 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -54,6 +54,7 @@ pub mod llm_obs; pub mod logs; pub mod logs_restriction; pub mod metrics; +pub mod model_lab; pub mod misc; pub mod monitors; pub mod ms_teams; diff --git a/src/commands/model_lab.rs b/src/commands/model_lab.rs new file mode 100644 index 00000000..77bf2ecf --- /dev/null +++ b/src/commands/model_lab.rs @@ -0,0 +1,520 @@ +use anyhow::Result; +use datadog_api_client::datadogV2::api_model_lab_api::{ + ListModelLabProjectsOptionalParams, ListModelLabRunArtifactsOptionalParams, + ListModelLabRunsOptionalParams, ModelLabAPIAPI, +}; +use datadog_api_client::datadogV2::model::{ + ModelLabFacetType, ModelLabProjectFacetType, ModelLabRunStatus, +}; + +use crate::config::Config; +use crate::formatter; + +fn make_api(cfg: &Config) -> ModelLabAPIAPI { + crate::make_api!(ModelLabAPIAPI, cfg) +} + +// --- Projects --- + +pub async fn projects_list( + cfg: &Config, + filter: Option, + filter_tags: Option, + sort: Option, + page_size: Option, + page_number: Option, +) -> Result<()> { + let api = make_api(cfg); + let mut params = ListModelLabProjectsOptionalParams::default(); + if let Some(f) = filter { + params = params.filter(f); + } + if let Some(t) = filter_tags { + params = params.filter_tags(t); + } + if let Some(s) = sort { + params = params.sort(s); + } + if let Some(n) = page_size { + params = params.page_size(n); + } + if let Some(n) = page_number { + params = params.page_number(n); + } + let resp = api + .list_model_lab_projects(params) + .await + .map_err(|e| anyhow::anyhow!("failed to list model lab projects: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn projects_get(cfg: &Config, project_id: i64) -> Result<()> { + let api = make_api(cfg); + let resp = api + .get_model_lab_project(project_id) + .await + .map_err(|e| anyhow::anyhow!("failed to get model lab project: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn projects_star(cfg: &Config, project_id: i64) -> Result<()> { + let api = make_api(cfg); + api.star_model_lab_project(project_id) + .await + .map_err(|e| anyhow::anyhow!("failed to star model lab project: {e:?}"))?; + println!("Project {project_id} starred."); + Ok(()) +} + +pub async fn projects_unstar(cfg: &Config, project_id: i64) -> Result<()> { + let api = make_api(cfg); + api.unstar_model_lab_project(project_id) + .await + .map_err(|e| anyhow::anyhow!("failed to unstar model lab project: {e:?}"))?; + println!("Project {project_id} unstarred."); + Ok(()) +} + +pub async fn projects_artifacts(cfg: &Config, project_id: i64) -> Result<()> { + let api = make_api(cfg); + let resp = api + .list_model_lab_project_artifacts(project_id) + .await + .map_err(|e| anyhow::anyhow!("failed to list model lab project artifacts: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn projects_facet_keys(cfg: &Config) -> Result<()> { + let api = make_api(cfg); + let resp = api + .list_model_lab_project_facet_keys() + .await + .map_err(|e| anyhow::anyhow!("failed to list model lab project facet keys: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn projects_facet_values( + cfg: &Config, + facet_type: &str, + facet_name: String, +) -> Result<()> { + let api = make_api(cfg); + let ft = match facet_type { + "tag" => ModelLabProjectFacetType::TAG, + other => anyhow::bail!( + "unknown facet type '{other}'; valid values: tag" + ), + }; + let resp = api + .list_model_lab_project_facet_values(ft, facet_name) + .await + .map_err(|e| anyhow::anyhow!("failed to list model lab project facet values: {e:?}"))?; + formatter::output(cfg, &resp) +} + +// --- Runs --- + +#[allow(clippy::too_many_arguments)] +pub async fn runs_list( + cfg: &Config, + filter: Option, + filter_project_id: Option, + filter_status: Option, + filter_tags: Option, + filter_params: Option, + filter_parent_run_id: Option, + pinned_first: bool, + include_pinned: bool, + sort: Option, + page_size: Option, + page_number: Option, +) -> Result<()> { + let api = make_api(cfg); + let mut params = ListModelLabRunsOptionalParams::default(); + if let Some(f) = filter { + params = params.filter(f); + } + if let Some(id) = filter_project_id { + params = params.filter_project_id(id); + } + if let Some(s) = filter_status { + let status = parse_run_status(&s)?; + params = params.filter_status(status); + } + if let Some(t) = filter_tags { + params = params.filter_tags(t); + } + if let Some(p) = filter_params { + params = params.filter_params(p); + } + if let Some(r) = filter_parent_run_id { + params = params.filter_parent_run_id(r); + } + if pinned_first { + params = params.pinned_first(true); + } + if include_pinned { + params = params.include_pinned(true); + } + if let Some(s) = sort { + params = params.sort(s); + } + if let Some(n) = page_size { + params = params.page_size(n); + } + if let Some(n) = page_number { + params = params.page_number(n); + } + let resp = api + .list_model_lab_runs(params) + .await + .map_err(|e| anyhow::anyhow!("failed to list model lab runs: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn runs_get(cfg: &Config, run_id: i64) -> Result<()> { + let api = make_api(cfg); + let resp = api + .get_model_lab_run(run_id) + .await + .map_err(|e| anyhow::anyhow!("failed to get model lab run: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn runs_delete(cfg: &Config, run_id: i64) -> Result<()> { + let api = make_api(cfg); + api.delete_model_lab_run(run_id) + .await + .map_err(|e| anyhow::anyhow!("failed to delete model lab run: {e:?}"))?; + println!("Run {run_id} deleted."); + Ok(()) +} + +pub async fn runs_pin(cfg: &Config, run_id: i64) -> Result<()> { + let api = make_api(cfg); + api.pin_model_lab_run(run_id) + .await + .map_err(|e| anyhow::anyhow!("failed to pin model lab run: {e:?}"))?; + println!("Run {run_id} pinned."); + Ok(()) +} + +pub async fn runs_unpin(cfg: &Config, run_id: i64) -> Result<()> { + let api = make_api(cfg); + api.unpin_model_lab_run(run_id) + .await + .map_err(|e| anyhow::anyhow!("failed to unpin model lab run: {e:?}"))?; + println!("Run {run_id} unpinned."); + Ok(()) +} + +pub async fn runs_artifacts(cfg: &Config, run_id: i64, path: Option) -> Result<()> { + let api = make_api(cfg); + let mut params = ListModelLabRunArtifactsOptionalParams::default(); + if let Some(p) = path { + params = params.path(p); + } + let resp = api + .list_model_lab_run_artifacts(run_id, params) + .await + .map_err(|e| anyhow::anyhow!("failed to list model lab run artifacts: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn runs_facet_keys(cfg: &Config, filter_project_id: i64) -> Result<()> { + let api = make_api(cfg); + let resp = api + .list_model_lab_run_facet_keys(filter_project_id) + .await + .map_err(|e| anyhow::anyhow!("failed to list model lab run facet keys: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn runs_facet_values( + cfg: &Config, + filter_project_id: i64, + facet_type: &str, + facet_name: String, +) -> Result<()> { + let api = make_api(cfg); + let ft = parse_facet_type(facet_type)?; + let resp = api + .list_model_lab_run_facet_values(filter_project_id, ft, facet_name) + .await + .map_err(|e| anyhow::anyhow!("failed to list model lab run facet values: {e:?}"))?; + formatter::output(cfg, &resp) +} + +fn parse_run_status(s: &str) -> Result { + match s { + "pending" => Ok(ModelLabRunStatus::PENDING), + "running" => Ok(ModelLabRunStatus::RUNNING), + "completed" => Ok(ModelLabRunStatus::COMPLETED), + "failed" => Ok(ModelLabRunStatus::FAILED), + "killed" => Ok(ModelLabRunStatus::KILLED), + "unresponsive" => Ok(ModelLabRunStatus::UNRESPONSIVE), + "paused" => Ok(ModelLabRunStatus::PAUSED), + other => anyhow::bail!( + "unknown status '{other}'; valid values: pending, running, completed, failed, killed, unresponsive, paused" + ), + } +} + +fn parse_facet_type(s: &str) -> Result { + match s { + "parameter" => Ok(ModelLabFacetType::PARAMETER), + "attribute" => Ok(ModelLabFacetType::ATTRIBUTE), + "tag" => Ok(ModelLabFacetType::TAG), + "metric" => Ok(ModelLabFacetType::METRIC), + other => anyhow::bail!( + "unknown facet type '{other}'; valid values: parameter, attribute, tag, metric" + ), + } +} + +#[cfg(test)] +mod tests { + use crate::test_support::*; + + // --- Projects --- + + #[tokio::test] + async fn test_model_lab_projects_list() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, r#"{"data":[],"meta":{"page":{"number":0,"size":10,"total":0}}}"#).await; + let result = super::projects_list(&cfg, None, None, None, None, None).await; + assert!(result.is_ok(), "projects list failed: {:?}", result.err()); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_projects_list_error() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + s.mock("GET", mockito::Matcher::Any) + .with_status(403) + .with_body(r#"{"errors":["Forbidden"]}"#) + .create_async() + .await; + let result = super::projects_list(&cfg, None, None, None, None, None).await; + assert!(result.is_err(), "projects list should fail on 403"); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_projects_get() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all( + &mut s, + r#"{"data":{"id":"1","type":"projects","attributes":{"artifact_storage_location":"s3://bucket","created_at":"2024-01-01T00:00:00Z","description":"test","is_starred":false,"name":"my-project","tags":[],"updated_at":"2024-01-01T00:00:00Z"}}}"#, + ) + .await; + let result = super::projects_get(&cfg, 1).await; + assert!(result.is_ok(), "projects get failed: {:?}", result.err()); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_projects_get_error() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + s.mock("GET", mockito::Matcher::Any) + .with_status(404) + .with_body(r#"{"errors":["not found"]}"#) + .create_async() + .await; + let result = super::projects_get(&cfg, 999).await; + assert!(result.is_err(), "projects get should fail on 404"); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_projects_star() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, "").await; + let result = super::projects_star(&cfg, 1).await; + assert!(result.is_ok(), "projects star failed: {:?}", result.err()); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_projects_unstar() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, "").await; + let result = super::projects_unstar(&cfg, 1).await; + assert!(result.is_ok(), "projects unstar failed: {:?}", result.err()); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_projects_artifacts() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, r#"{"data":{"id":"1","type":"project_files","attributes":{"files":[]}}}"#).await; + let result = super::projects_artifacts(&cfg, 1).await; + assert!( + result.is_ok(), + "projects artifacts failed: {:?}", + result.err() + ); + cleanup_env(); + } + + // --- Runs --- + + #[tokio::test] + async fn test_model_lab_runs_list() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, r#"{"data":[],"meta":{"page":{"number":0,"size":10,"total":0}}}"#).await; + let result = + super::runs_list(&cfg, None, None, None, None, None, None, false, false, None, None, None) + .await; + assert!(result.is_ok(), "runs list failed: {:?}", result.err()); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_runs_list_invalid_status() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, r#"{"data":[]}"#).await; + let result = super::runs_list( + &cfg, + None, + None, + Some("bogus".to_string()), + None, + None, + None, + false, + false, + None, + None, + None, + ) + .await; + assert!(result.is_err(), "invalid status should fail"); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_runs_get() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all( + &mut s, + r#"{"data":{"id":"42","type":"runs","attributes":{"created_at":"2024-01-01T00:00:00Z","descendant_match":false,"description":"test run","has_children":false,"is_pinned":false,"metric_summaries":[],"mlflow_artifact_location":"s3://bucket/run","name":"run-1","params":null,"project_id":1,"started_at":"2024-01-01T00:00:00Z","status":"completed","tags":[],"updated_at":"2024-01-01T00:00:00Z"}}}"#, + ) + .await; + let result = super::runs_get(&cfg, 42).await; + assert!(result.is_ok(), "runs get failed: {:?}", result.err()); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_runs_get_error() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + s.mock("GET", mockito::Matcher::Any) + .with_status(404) + .with_body(r#"{"errors":["not found"]}"#) + .create_async() + .await; + let result = super::runs_get(&cfg, 999).await; + assert!(result.is_err(), "runs get should fail on 404"); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_runs_delete() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, "").await; + let result = super::runs_delete(&cfg, 42).await; + assert!(result.is_ok(), "runs delete failed: {:?}", result.err()); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_runs_delete_error() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + s.mock("DELETE", mockito::Matcher::Any) + .with_status(404) + .with_body(r#"{"errors":["not found"]}"#) + .create_async() + .await; + let result = super::runs_delete(&cfg, 999).await; + assert!(result.is_err(), "runs delete should fail on 404"); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_runs_pin() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, "").await; + let result = super::runs_pin(&cfg, 42).await; + assert!(result.is_ok(), "runs pin failed: {:?}", result.err()); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_runs_unpin() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, "").await; + let result = super::runs_unpin(&cfg, 42).await; + assert!(result.is_ok(), "runs unpin failed: {:?}", result.err()); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_runs_artifacts() { + let _lock = lock_env().await; + let mut s = mockito::Server::new_async().await; + let cfg = test_config(&s.url()); + mock_all(&mut s, r#"{"data":{"id":"42","type":"artifacts","attributes":{"files":[],"path_in_project":""}}}"#).await; + let result = super::runs_artifacts(&cfg, 42, None).await; + assert!( + result.is_ok(), + "runs artifacts failed: {:?}", + result.err() + ); + cleanup_env(); + } + + #[tokio::test] + async fn test_model_lab_parse_run_status_valid() { + for s in &[ + "pending", "running", "completed", "failed", "killed", "unresponsive", "paused", + ] { + assert!(super::parse_run_status(s).is_ok(), "expected '{s}' to be valid"); + } + } + + #[tokio::test] + async fn test_model_lab_parse_run_status_invalid() { + assert!(super::parse_run_status("bogus").is_err()); + } +} diff --git a/src/config.rs b/src/config.rs index 4e3d38a2..d82aae22 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1936,6 +1936,7 @@ mod tests { auto_approve: false, agent_mode: false, read_only: false, + jq: None, }; super::apply_org_override(&mut cfg, "unknown-org".into()).unwrap(); diff --git a/src/main.rs b/src/main.rs index f95f1ad0..75ea065e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1879,6 +1879,30 @@ enum Commands { #[command(subcommand)] action: MiscActions, }, + /// Explore Model Lab projects and runs + /// + /// Browse ML experiment projects and their training runs in Datadog Model Lab. + /// + /// SUBCOMMANDS: + /// projects Manage Model Lab projects + /// runs Manage Model Lab runs + /// + /// EXAMPLES: + /// pup model-lab projects list + /// pup model-lab projects get 123 + /// pup model-lab projects star 123 + /// pup model-lab runs list --filter-project-id 123 + /// pup model-lab runs get 456 + /// pup model-lab runs delete 456 + /// pup model-lab runs pin 456 + /// + /// AUTHENTICATION: + /// Requires either OAuth2 authentication or API keys. + #[command(name = "model-lab", verbatim_doc_comment)] + ModelLab { + #[command(subcommand)] + action: ModelLabActions, + }, /// Manage monitors /// /// Manage Datadog monitors for alerting and notifications. @@ -8034,6 +8058,119 @@ enum MiscActions { Status, } +// ---- Model Lab ---- +#[derive(Subcommand)] +enum ModelLabActions { + /// Manage Model Lab projects + Projects { + #[command(subcommand)] + action: ModelLabProjectActions, + }, + /// Manage Model Lab runs + Runs { + #[command(subcommand)] + action: ModelLabRunActions, + }, +} + +#[derive(Subcommand)] +enum ModelLabProjectActions { + /// List Model Lab projects + List { + #[arg(long, help = "Filter query string")] + filter: Option, + #[arg(long, help = "Filter by tags (comma-separated)")] + filter_tags: Option, + #[arg(long, help = "Sort field (e.g. name, created_at)")] + sort: Option, + #[arg(long)] + page_size: Option, + #[arg(long, help = "Page number (0-indexed)")] + page_number: Option, + }, + /// Get a Model Lab project by ID + Get { project_id: i64 }, + /// Star a Model Lab project + Star { project_id: i64 }, + /// Unstar a Model Lab project + Unstar { project_id: i64 }, + /// List artifacts for a Model Lab project + Artifacts { project_id: i64 }, + /// List facet keys for Model Lab projects + #[command(name = "facet-keys")] + FacetKeys, + /// List facet values for a Model Lab project facet + #[command(name = "facet-values")] + FacetValues { + #[arg(help = "Facet type (tag)")] + facet_type: String, + #[arg(help = "Facet name")] + facet_name: String, + }, +} + +#[derive(Subcommand)] +enum ModelLabRunActions { + /// List Model Lab runs + List { + #[arg(long, help = "Filter query string")] + filter: Option, + #[arg(long, help = "Filter by project ID")] + filter_project_id: Option, + #[arg( + long, + help = "Filter by status: pending, running, completed, failed, killed, unresponsive, paused" + )] + filter_status: Option, + #[arg(long, help = "Filter by tags (comma-separated)")] + filter_tags: Option, + #[arg(long, help = "Filter by params")] + filter_params: Option, + #[arg(long, help = "Filter by parent run ID")] + filter_parent_run_id: Option, + #[arg(long, default_value_t = false, help = "Show pinned runs first")] + pinned_first: bool, + #[arg(long, default_value_t = false, help = "Include pinned runs")] + include_pinned: bool, + #[arg(long, help = "Sort field")] + sort: Option, + #[arg(long)] + page_size: Option, + #[arg(long, help = "Page number (0-indexed)")] + page_number: Option, + }, + /// Get a Model Lab run by ID + Get { run_id: i64 }, + /// Delete a Model Lab run + Delete { run_id: i64 }, + /// Pin a Model Lab run + Pin { run_id: i64 }, + /// Unpin a Model Lab run + Unpin { run_id: i64 }, + /// List artifacts for a Model Lab run + Artifacts { + run_id: i64, + #[arg(long, help = "Filter by artifact path prefix")] + path: Option, + }, + /// List facet keys for Model Lab runs + #[command(name = "facet-keys")] + FacetKeys { + #[arg(long, help = "Project ID to scope facet keys")] + filter_project_id: i64, + }, + /// List facet values for a Model Lab run facet + #[command(name = "facet-values")] + FacetValues { + #[arg(long, help = "Project ID to scope facet values")] + filter_project_id: i64, + #[arg(help = "Facet type (parameter, attribute, tag, metric)")] + facet_type: String, + #[arg(help = "Facet name")] + facet_name: String, + }, +} + // ---- APM ---- #[derive(Subcommand)] enum ApmActions { @@ -14169,6 +14306,115 @@ async fn main_inner() -> anyhow::Result<()> { MiscActions::Status => commands::misc::status(&cfg).await?, } } + // --- Model Lab --- + Commands::ModelLab { action } => { + cfg.validate_auth()?; + match action { + ModelLabActions::Projects { action } => match action { + ModelLabProjectActions::List { + filter, + filter_tags, + sort, + page_size, + page_number, + } => { + commands::model_lab::projects_list( + &cfg, + filter, + filter_tags, + sort, + page_size, + page_number, + ) + .await?; + } + ModelLabProjectActions::Get { project_id } => { + commands::model_lab::projects_get(&cfg, project_id).await?; + } + ModelLabProjectActions::Star { project_id } => { + commands::model_lab::projects_star(&cfg, project_id).await?; + } + ModelLabProjectActions::Unstar { project_id } => { + commands::model_lab::projects_unstar(&cfg, project_id).await?; + } + ModelLabProjectActions::Artifacts { project_id } => { + commands::model_lab::projects_artifacts(&cfg, project_id).await?; + } + ModelLabProjectActions::FacetKeys => { + commands::model_lab::projects_facet_keys(&cfg).await?; + } + ModelLabProjectActions::FacetValues { + facet_type, + facet_name, + } => { + commands::model_lab::projects_facet_values(&cfg, &facet_type, facet_name) + .await?; + } + }, + ModelLabActions::Runs { action } => match action { + ModelLabRunActions::List { + filter, + filter_project_id, + filter_status, + filter_tags, + filter_params, + filter_parent_run_id, + pinned_first, + include_pinned, + sort, + page_size, + page_number, + } => { + commands::model_lab::runs_list( + &cfg, + filter, + filter_project_id, + filter_status, + filter_tags, + filter_params, + filter_parent_run_id, + pinned_first, + include_pinned, + sort, + page_size, + page_number, + ) + .await?; + } + ModelLabRunActions::Get { run_id } => { + commands::model_lab::runs_get(&cfg, run_id).await?; + } + ModelLabRunActions::Delete { run_id } => { + commands::model_lab::runs_delete(&cfg, run_id).await?; + } + ModelLabRunActions::Pin { run_id } => { + commands::model_lab::runs_pin(&cfg, run_id).await?; + } + ModelLabRunActions::Unpin { run_id } => { + commands::model_lab::runs_unpin(&cfg, run_id).await?; + } + ModelLabRunActions::Artifacts { run_id, path } => { + commands::model_lab::runs_artifacts(&cfg, run_id, path).await?; + } + ModelLabRunActions::FacetKeys { filter_project_id } => { + commands::model_lab::runs_facet_keys(&cfg, filter_project_id).await?; + } + ModelLabRunActions::FacetValues { + filter_project_id, + facet_type, + facet_name, + } => { + commands::model_lab::runs_facet_values( + &cfg, + filter_project_id, + &facet_type, + facet_name, + ) + .await?; + } + }, + } + } // --- APM --- Commands::Apm { action } => { cfg.validate_auth()?; From 7ec327e864eea972e72b723a75900339f4cf66a4 Mon Sep 17 00:00:00 2001 From: jack-edmonds-dd Date: Thu, 18 Jun 2026 16:55:46 -0400 Subject: [PATCH 2/2] fmt --- src/commands/mod.rs | 2 +- src/commands/model_lab.rs | 48 ++++++++++++++++++++++++++------------- 2 files changed, 33 insertions(+), 17 deletions(-) diff --git a/src/commands/mod.rs b/src/commands/mod.rs index a28ce14c..0eacb306 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -54,8 +54,8 @@ pub mod llm_obs; pub mod logs; pub mod logs_restriction; pub mod metrics; -pub mod model_lab; pub mod misc; +pub mod model_lab; pub mod monitors; pub mod ms_teams; pub mod network; diff --git a/src/commands/model_lab.rs b/src/commands/model_lab.rs index 77bf2ecf..04379d68 100644 --- a/src/commands/model_lab.rs +++ b/src/commands/model_lab.rs @@ -101,9 +101,7 @@ pub async fn projects_facet_values( let api = make_api(cfg); let ft = match facet_type { "tag" => ModelLabProjectFacetType::TAG, - other => anyhow::bail!( - "unknown facet type '{other}'; valid values: tag" - ), + other => anyhow::bail!("unknown facet type '{other}'; valid values: tag"), }; let resp = api .list_model_lab_project_facet_values(ft, facet_name) @@ -283,7 +281,11 @@ mod tests { let _lock = lock_env().await; let mut s = mockito::Server::new_async().await; let cfg = test_config(&s.url()); - mock_all(&mut s, r#"{"data":[],"meta":{"page":{"number":0,"size":10,"total":0}}}"#).await; + mock_all( + &mut s, + r#"{"data":[],"meta":{"page":{"number":0,"size":10,"total":0}}}"#, + ) + .await; let result = super::projects_list(&cfg, None, None, None, None, None).await; assert!(result.is_ok(), "projects list failed: {:?}", result.err()); cleanup_env(); @@ -361,7 +363,11 @@ mod tests { let _lock = lock_env().await; let mut s = mockito::Server::new_async().await; let cfg = test_config(&s.url()); - mock_all(&mut s, r#"{"data":{"id":"1","type":"project_files","attributes":{"files":[]}}}"#).await; + mock_all( + &mut s, + r#"{"data":{"id":"1","type":"project_files","attributes":{"files":[]}}}"#, + ) + .await; let result = super::projects_artifacts(&cfg, 1).await; assert!( result.is_ok(), @@ -378,10 +384,15 @@ mod tests { let _lock = lock_env().await; let mut s = mockito::Server::new_async().await; let cfg = test_config(&s.url()); - mock_all(&mut s, r#"{"data":[],"meta":{"page":{"number":0,"size":10,"total":0}}}"#).await; - let result = - super::runs_list(&cfg, None, None, None, None, None, None, false, false, None, None, None) - .await; + mock_all( + &mut s, + r#"{"data":[],"meta":{"page":{"number":0,"size":10,"total":0}}}"#, + ) + .await; + let result = super::runs_list( + &cfg, None, None, None, None, None, None, false, false, None, None, None, + ) + .await; assert!(result.is_ok(), "runs list failed: {:?}", result.err()); cleanup_env(); } @@ -496,20 +507,25 @@ mod tests { let cfg = test_config(&s.url()); mock_all(&mut s, r#"{"data":{"id":"42","type":"artifacts","attributes":{"files":[],"path_in_project":""}}}"#).await; let result = super::runs_artifacts(&cfg, 42, None).await; - assert!( - result.is_ok(), - "runs artifacts failed: {:?}", - result.err() - ); + assert!(result.is_ok(), "runs artifacts failed: {:?}", result.err()); cleanup_env(); } #[tokio::test] async fn test_model_lab_parse_run_status_valid() { for s in &[ - "pending", "running", "completed", "failed", "killed", "unresponsive", "paused", + "pending", + "running", + "completed", + "failed", + "killed", + "unresponsive", + "paused", ] { - assert!(super::parse_run_status(s).is_ok(), "expected '{s}' to be valid"); + assert!( + super::parse_run_status(s).is_ok(), + "expected '{s}' to be valid" + ); } }