diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index 0948138..6f0dbd7 100644 --- a/docs/COMMANDS.md +++ b/docs/COMMANDS.md @@ -66,7 +66,7 @@ pup [options] # Nested commands | data-deletion | requests (list, create, cancel) | src/commands/data_deletion.rs | ✅ | | data-governance | scanner-rules (list) | src/commands/data_governance.rs | ✅ | | obs-pipelines | list, get, create, update, delete, validate | src/commands/obs_pipelines.rs | ✅ | -| llm-obs | projects (create, list), experiments (create, list, update, delete, summary, events (list, get), metric-values, dimension-values), datasets (create, list), spans (search) | src/commands/llm_obs.rs | ✅ | +| llm-obs | projects (create, list), experiments (create, list, update, delete, summary, events (list, get), metric-values, dimension-values), datasets (create, list, batch-update, clone, restore), spans (search) | src/commands/llm_obs.rs | ✅ | | reference-tables | list, get, create, batch-query | src/commands/reference_tables.rs | ✅ | | network | flows list, devices (list, get, interfaces, tags), interfaces (list, update) | src/commands/network.rs | ✅ | | cloud | aws, gcp, azure, oci | src/commands/cloud.rs | ✅ | @@ -243,7 +243,7 @@ Available on all commands: ### v0.28.0 — New Command Groups and Full Pipeline Implementation -- ✅ **llm-obs** (new) — LLM Observability: projects (create, list), experiments (create, list, update, delete, summary, events (list, get), metric-values, dimension-values), datasets (create, list), spans (search) +- ✅ **llm-obs** (new) — LLM Observability: projects (create, list), experiments (create, list, update, delete, summary, events (list, get), metric-values, dimension-values), datasets (create, list, batch-update, clone, restore), spans (search) - ✅ **reference-tables** (new) — Reference table management (list, get, create, batch-query) - ✅ **obs-pipelines** (upgraded from placeholder) — Full CRUD: list, get, create, update, delete, validate - **costs** — Added cloud cost configs: `aws-config`, `azure-config`, `gcp-config` (list, get, create, delete each) diff --git a/src/client.rs b/src/client.rs index 42cc95d..fefb09f 100644 --- a/src/client.rs +++ b/src/client.rs @@ -350,7 +350,7 @@ static UNSTABLE_OPS: &[&str] = &[ "v2.delete_aws_cloud_auth_persona_mapping", "v2.get_aws_cloud_auth_persona_mapping", "v2.list_aws_cloud_auth_persona_mappings", - // LLM Observability (18) + // LLM Observability (21) "v2.create_llm_obs_project", "v2.list_llm_obs_projects", "v2.create_llm_obs_experiment", @@ -359,6 +359,9 @@ static UNSTABLE_OPS: &[&str] = &[ "v2.delete_llm_obs_experiments", "v2.create_llm_obs_dataset", "v2.list_llm_obs_datasets", + "v2.batch_update_llm_obs_dataset", + "v2.clone_llm_obs_dataset", + "v2.restore_llm_obs_dataset_version", "v2.create_llm_obs_annotation_queue", "v2.list_llm_obs_annotation_queues", "v2.update_llm_obs_annotation_queue", @@ -1273,7 +1276,7 @@ mod tests { #[test] fn test_unstable_ops_count() { - assert_eq!(UNSTABLE_OPS.len(), 162); + assert_eq!(UNSTABLE_OPS.len(), 165); } #[test] diff --git a/src/commands/llm_obs.rs b/src/commands/llm_obs.rs index 6532c62..af05a95 100644 --- a/src/commands/llm_obs.rs +++ b/src/commands/llm_obs.rs @@ -4,9 +4,11 @@ use datadog_api_client::datadogV2::api_llm_observability::{ }; use datadog_api_client::datadogV2::model::{ LLMObsAnnotationQueueInteractionsRequest, LLMObsAnnotationQueueRequest, - LLMObsAnnotationQueueUpdateRequest, LLMObsCustomEvalConfigUpdateRequest, LLMObsDatasetRequest, - LLMObsDeleteAnnotationQueueInteractionsRequest, LLMObsDeleteExperimentsRequest, - LLMObsExperimentRequest, LLMObsExperimentUpdateRequest, LLMObsProjectRequest, + LLMObsAnnotationQueueUpdateRequest, LLMObsCustomEvalConfigUpdateRequest, + LLMObsDatasetBatchUpdateRequest, LLMObsDatasetCloneRequest, LLMObsDatasetRequest, + LLMObsDatasetRestoreVersionRequest, LLMObsDeleteAnnotationQueueInteractionsRequest, + LLMObsDeleteExperimentsRequest, LLMObsExperimentRequest, LLMObsExperimentUpdateRequest, + LLMObsProjectRequest, }; use crate::client; @@ -108,6 +110,51 @@ pub async fn datasets_list(cfg: &Config, project_id: &str) -> Result<()> { formatter::output(cfg, &resp) } +pub async fn datasets_batch_update( + cfg: &Config, + project_id: &str, + dataset_id: &str, + file: &str, +) -> Result<()> { + let body: LLMObsDatasetBatchUpdateRequest = util::read_json_file(file)?; + let api = make_api(cfg); + let resp = api + .batch_update_llm_obs_dataset(project_id.to_string(), dataset_id.to_string(), body) + .await + .map_err(|e| anyhow::anyhow!("failed to batch update dataset records: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn datasets_clone( + cfg: &Config, + project_id: &str, + dataset_id: &str, + file: &str, +) -> Result<()> { + let body: LLMObsDatasetCloneRequest = util::read_json_file(file)?; + let api = make_api(cfg); + let resp = api + .clone_llm_obs_dataset(project_id.to_string(), dataset_id.to_string(), body) + .await + .map_err(|e| anyhow::anyhow!("failed to clone dataset: {e:?}"))?; + formatter::output(cfg, &resp) +} + +pub async fn datasets_restore( + cfg: &Config, + project_id: &str, + dataset_id: &str, + file: &str, +) -> Result<()> { + let body: LLMObsDatasetRestoreVersionRequest = util::read_json_file(file)?; + let api = make_api(cfg); + api.restore_llm_obs_dataset_version(project_id.to_string(), dataset_id.to_string(), body) + .await + .map_err(|e| anyhow::anyhow!("failed to restore dataset version: {e:?}"))?; + println!("Dataset {dataset_id} restored."); + Ok(()) +} + // ---- Experiment analytics (no typed equivalent — unstable MCP endpoints) ---- pub async fn experiments_summary(cfg: &Config, experiment_id: &str) -> Result<()> { @@ -2637,4 +2684,165 @@ mod tests { assert!(result.is_ok(), "spans_search failed: {:?}", result.err()); cleanup_env(); } + + // ---- datasets_batch_update ---- + + #[tokio::test] + async fn test_llm_obs_datasets_batch_update() { + let _lock = lock_env().await; + std::env::set_var("DD_TOKEN_STORAGE", "file"); + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let tmp = write_temp_json( + "pup_test_ds_batch_update.json", + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"insert_records":[],"delete_records":[]}}}"#, + ); + let resp_body = r#"{"data":[]}"#; + let _mock = mock_any(&mut server, "POST", resp_body).await; + + let result = + super::datasets_batch_update(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await; + assert!( + result.is_ok(), + "datasets_batch_update failed: {:?}", + result.err() + ); + let _ = std::fs::remove_file(tmp); + cleanup_env(); + std::env::remove_var("DD_TOKEN_STORAGE"); + } + + #[tokio::test] + async fn test_llm_obs_datasets_batch_update_400() { + let _lock = lock_env().await; + std::env::set_var("DD_TOKEN_STORAGE", "file"); + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let tmp = write_temp_json( + "pup_test_ds_batch_update_400.json", + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"insert_records":[],"delete_records":[]}}}"#, + ); + let _mock = server + .mock("POST", mockito::Matcher::Any) + .match_query(mockito::Matcher::Any) + .with_status(400) + .with_header("content-type", "application/json") + .with_body(r#"{"errors":["bad request"]}"#) + .create_async() + .await; + + let result = + super::datasets_batch_update(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await; + assert!(result.is_err(), "should fail on 400"); + let _ = std::fs::remove_file(tmp); + cleanup_env(); + std::env::remove_var("DD_TOKEN_STORAGE"); + } + + // ---- datasets_clone ---- + + #[tokio::test] + async fn test_llm_obs_datasets_clone() { + let _lock = lock_env().await; + std::env::set_var("DD_TOKEN_STORAGE", "file"); + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let tmp = write_temp_json( + "pup_test_ds_clone.json", + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"name":"cloned-dataset"}}}"#, + ); + let resp_body = r#"{"data":{"id":"ds-2","type":"datasets","attributes":{"name":"cloned-dataset","description":null,"metadata":null,"created_at":"2024-01-01T00:00:00Z","updated_at":"2024-01-01T00:00:00Z","current_version":1}}}"#; + let _mock = mock_any(&mut server, "POST", resp_body).await; + + let result = super::datasets_clone(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await; + assert!(result.is_ok(), "datasets_clone failed: {:?}", result.err()); + let _ = std::fs::remove_file(tmp); + cleanup_env(); + std::env::remove_var("DD_TOKEN_STORAGE"); + } + + #[tokio::test] + async fn test_llm_obs_datasets_clone_404() { + let _lock = lock_env().await; + std::env::set_var("DD_TOKEN_STORAGE", "file"); + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let tmp = write_temp_json( + "pup_test_ds_clone_404.json", + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"name":"cloned-dataset"}}}"#, + ); + let _mock = server + .mock("POST", mockito::Matcher::Any) + .match_query(mockito::Matcher::Any) + .with_status(404) + .with_header("content-type", "application/json") + .with_body(r#"{"errors":["not found"]}"#) + .create_async() + .await; + + let result = + super::datasets_clone(&cfg, "proj-1", "ds-missing", tmp.to_str().unwrap()).await; + assert!(result.is_err(), "should fail on 404"); + let _ = std::fs::remove_file(tmp); + cleanup_env(); + std::env::remove_var("DD_TOKEN_STORAGE"); + } + + // ---- datasets_restore ---- + + #[tokio::test] + async fn test_llm_obs_datasets_restore() { + let _lock = lock_env().await; + std::env::set_var("DD_TOKEN_STORAGE", "file"); + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let tmp = write_temp_json( + "pup_test_ds_restore.json", + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"dataset_version":2}}}"#, + ); + let resp_body = r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"name":"my-dataset","description":null,"metadata":null,"created_at":"2024-01-01T00:00:00Z","updated_at":"2024-01-01T00:00:00Z","current_version":2}}}"#; + let _mock = mock_any(&mut server, "POST", resp_body).await; + + let result = super::datasets_restore(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await; + assert!( + result.is_ok(), + "datasets_restore failed: {:?}", + result.err() + ); + let _ = std::fs::remove_file(tmp); + cleanup_env(); + std::env::remove_var("DD_TOKEN_STORAGE"); + } + + #[tokio::test] + async fn test_llm_obs_datasets_restore_400() { + let _lock = lock_env().await; + std::env::set_var("DD_TOKEN_STORAGE", "file"); + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let tmp = write_temp_json( + "pup_test_ds_restore_400.json", + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"dataset_version":99}}}"#, + ); + let _mock = server + .mock("POST", mockito::Matcher::Any) + .match_query(mockito::Matcher::Any) + .with_status(400) + .with_header("content-type", "application/json") + .with_body(r#"{"errors":["invalid version"]}"#) + .create_async() + .await; + + let result = super::datasets_restore(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await; + assert!(result.is_err(), "should fail on 400"); + let _ = std::fs::remove_file(tmp); + cleanup_env(); + std::env::remove_var("DD_TOKEN_STORAGE"); + } } diff --git a/src/main.rs b/src/main.rs index 8fbdede..fc8ceac 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8682,6 +8682,33 @@ enum LlmObsDatasetsActions { #[arg(long, help = "Project ID (required)")] project_id: String, }, + /// Batch insert, update, and delete records in a dataset + BatchUpdate { + #[arg(long, help = "Project ID (required)")] + project_id: String, + #[arg(long, help = "Dataset ID (required)")] + dataset_id: String, + #[arg(long, help = "JSON file with batch update body (required)")] + file: String, + }, + /// Clone a dataset into a new dataset + Clone { + #[arg(long, help = "Project ID (required)")] + project_id: String, + #[arg(long, help = "Dataset ID to clone (required)")] + dataset_id: String, + #[arg(long, help = "JSON file with clone body (required)")] + file: String, + }, + /// Restore a dataset to a previous version + Restore { + #[arg(long, help = "Project ID (required)")] + project_id: String, + #[arg(long, help = "Dataset ID (required)")] + dataset_id: String, + #[arg(long, help = "JSON file with restore version body (required)")] + file: String, + }, } #[derive(Subcommand)] @@ -14661,6 +14688,35 @@ async fn main_inner() -> anyhow::Result<()> { LlmObsDatasetsActions::List { project_id } => { commands::llm_obs::datasets_list(&cfg, &project_id).await?; } + LlmObsDatasetsActions::BatchUpdate { + project_id, + dataset_id, + file, + } => { + commands::llm_obs::datasets_batch_update( + &cfg, + &project_id, + &dataset_id, + &file, + ) + .await?; + } + LlmObsDatasetsActions::Clone { + project_id, + dataset_id, + file, + } => { + commands::llm_obs::datasets_clone(&cfg, &project_id, &dataset_id, &file) + .await?; + } + LlmObsDatasetsActions::Restore { + project_id, + dataset_id, + file, + } => { + commands::llm_obs::datasets_restore(&cfg, &project_id, &dataset_id, &file) + .await?; + } }, LlmObsActions::Spans { action } => match action { LlmObsSpansActions::Search {