From 0effe439e050b10ee987991c8b875561ac698534 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 29 May 2026 11:53:34 +0000 Subject: [PATCH 1/3] feat(llm-obs): add dataset batch-update, clone, restore subcommands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wrap three new SDK endpoints introduced in datadog-api-client-rust PR #1655 (available at rev d4954b11 picked up by chore/upgrade-dd-sdk-to-master): - `pup llm-obs datasets batch-update` — batch insert/update/delete records via LLMObsDatasetBatchUpdateRequest / batch_update_llm_obs_dataset - `pup llm-obs datasets clone` — clone a dataset into a new dataset via LLMObsDatasetCloneRequest / clone_llm_obs_dataset - `pup llm-obs datasets restore` — restore a dataset to a previous version via LLMObsDatasetRestoreVersionRequest / restore_llm_obs_dataset_version Changes: - src/commands/llm_obs.rs: add datasets_batch_update, datasets_clone, datasets_restore functions; add 6 tests (happy-path + error for each) - src/main.rs: add BatchUpdate/Clone/Restore variants to LlmObsDatasetsActions and dispatch arms - src/client.rs: add 3 new unstable op IDs; update count assertion 162 → 165 - docs/COMMANDS.md: update llm-obs datasets command listing Co-Authored-By: Claude https://claude.ai/code/session_01MaHfvEkY66q99gwcPZi4Xc --- docs/COMMANDS.md | 4 +- src/client.rs | 7 +- src/commands/llm_obs.rs | 214 +++++++++++++++++++++++++++++++++++++++- src/main.rs | 56 +++++++++++ 4 files changed, 274 insertions(+), 7 deletions(-) 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..216463b 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); + let resp = 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:?}"))?; + formatter::output(cfg, &resp) +} + // ---- 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":{"type":"dataset_batch_update","attributes":{"upsert":[],"delete":[]}}}"#, + ); + let resp_body = r#"{"data":{"type":"dataset_records_mutation","attributes":{"upserted_ids":[],"deleted_ids":[]}}}"#; + 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":{"type":"dataset_batch_update","attributes":{"upsert":[],"delete":[]}}}"#, + ); + 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":{"type":"dataset_clone","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":{"type":"dataset_clone","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":{"type":"dataset_restore","attributes":{"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":{"type":"dataset_restore","attributes":{"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 { From 7c2a872d97d0943c26e56f33e77920f90e83d557 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 29 May 2026 12:19:55 +0000 Subject: [PATCH 2/3] fix(llm-obs): drop unit-valued binding in datasets_restore restore_llm_obs_dataset_version returns no body; binding its unit result tripped clippy::let_unit_value (-D warnings) and broke the Check/Test and Windows cross-compile jobs. Follow the delete pattern: call, then print a confirmation. Co-Authored-By: Claude --- src/commands/llm_obs.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/commands/llm_obs.rs b/src/commands/llm_obs.rs index 216463b..40004f7 100644 --- a/src/commands/llm_obs.rs +++ b/src/commands/llm_obs.rs @@ -148,11 +148,11 @@ pub async fn datasets_restore( ) -> Result<()> { let body: LLMObsDatasetRestoreVersionRequest = util::read_json_file(file)?; let api = make_api(cfg); - let resp = api - .restore_llm_obs_dataset_version(project_id.to_string(), dataset_id.to_string(), body) + 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:?}"))?; - formatter::output(cfg, &resp) + println!("Dataset {dataset_id} restored."); + Ok(()) } // ---- Experiment analytics (no typed equivalent — unstable MCP endpoints) ---- From 29f11a3616cb48be1d09482988ea3af58517b7e3 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 29 May 2026 12:44:44 +0000 Subject: [PATCH 3/3] test(llm-obs): fix dataset request bodies to match SDK schema The batch-update/clone/restore tests built request JSON with stale field names (type "dataset_*", missing required id, wrong attribute keys), so util::read_json_file failed before the API call and the happy-path asserts panicked. Align bodies with the real SDK models: - data.id required; data.type is "datasets" - batch-update attrs use insert_records/delete_records - restore attr is dataset_version (not version) - batch-update response data is an array (LLMObsDatasetRecordsMutationResponse) Co-Authored-By: Claude --- src/commands/llm_obs.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/commands/llm_obs.rs b/src/commands/llm_obs.rs index 40004f7..af05a95 100644 --- a/src/commands/llm_obs.rs +++ b/src/commands/llm_obs.rs @@ -2696,9 +2696,9 @@ mod tests { let tmp = write_temp_json( "pup_test_ds_batch_update.json", - r#"{"data":{"type":"dataset_batch_update","attributes":{"upsert":[],"delete":[]}}}"#, + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"insert_records":[],"delete_records":[]}}}"#, ); - let resp_body = r#"{"data":{"type":"dataset_records_mutation","attributes":{"upserted_ids":[],"deleted_ids":[]}}}"#; + let resp_body = r#"{"data":[]}"#; let _mock = mock_any(&mut server, "POST", resp_body).await; let result = @@ -2722,7 +2722,7 @@ mod tests { let tmp = write_temp_json( "pup_test_ds_batch_update_400.json", - r#"{"data":{"type":"dataset_batch_update","attributes":{"upsert":[],"delete":[]}}}"#, + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"insert_records":[],"delete_records":[]}}}"#, ); let _mock = server .mock("POST", mockito::Matcher::Any) @@ -2752,7 +2752,7 @@ mod tests { let tmp = write_temp_json( "pup_test_ds_clone.json", - r#"{"data":{"type":"dataset_clone","attributes":{"name":"cloned-dataset"}}}"#, + 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; @@ -2773,7 +2773,7 @@ mod tests { let tmp = write_temp_json( "pup_test_ds_clone_404.json", - r#"{"data":{"type":"dataset_clone","attributes":{"name":"cloned-dataset"}}}"#, + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"name":"cloned-dataset"}}}"#, ); let _mock = server .mock("POST", mockito::Matcher::Any) @@ -2803,7 +2803,7 @@ mod tests { let tmp = write_temp_json( "pup_test_ds_restore.json", - r#"{"data":{"type":"dataset_restore","attributes":{"version":2}}}"#, + 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; @@ -2828,7 +2828,7 @@ mod tests { let tmp = write_temp_json( "pup_test_ds_restore_400.json", - r#"{"data":{"type":"dataset_restore","attributes":{"version":99}}}"#, + r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"dataset_version":99}}}"#, ); let _mock = server .mock("POST", mockito::Matcher::Any)