Merge pull request #728 from gsvigruha/gergely.svigruha/annotation-pup-cli

feat(llm-obs): add annotation queue label schema and annotation commands

Author: platinummonkey

docs/COMMANDS.md

@@ -66,7 +66,7 @@ pup <domain> <subgroup> <action> [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, submit), metric-values, dimension-values), datasets (create, list, batch-update, clone, restore, records, records-add, records-all, records-full), spans (search), patterns (configs (list, get), runs (list, status), topics, topics-with-points, points), agent-insights (list, get, update-status, submit-feedback), model-pricing | src/commands/llm_obs.rs | ✅ |
+| llm-obs | projects (create, list), experiments (create, list, update, delete, summary, events (list, get, submit), metric-values, dimension-values), datasets (create, list, batch-update, clone, restore, records, records-add, records-all, records-full), spans (search), patterns (configs (list, get), runs (list, status), topics, topics-with-points, points), agent-insights (list, get, update-status, submit-feedback), annotation-queues (create, list, update, delete, interactions (add, delete, list), schema (get, update), annotations (upsert, delete)), model-pricing | 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 | ✅ |

src/client.rs

@@ -301,7 +301,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 (21)
+    // LLM Observability (25)
     "v2.create_llm_obs_project",
     "v2.list_llm_obs_projects",
     "v2.create_llm_obs_experiment",
@@ -320,6 +320,10 @@ static UNSTABLE_OPS: &[&str] = &[
     "v2.create_llm_obs_annotation_queue_interactions",
     "v2.delete_llm_obs_annotation_queue_interactions",
     "v2.get_llm_obs_annotated_interactions",
+    "v2.get_llm_obs_annotation_queue_label_schema",
+    "v2.update_llm_obs_annotation_queue_label_schema",
+    "v2.upsert_llm_obs_annotations",
+    "v2.delete_llm_obs_annotations",
     "v2.get_llm_obs_custom_eval_config",
     "v2.update_llm_obs_custom_eval_config",
     "v2.delete_llm_obs_custom_eval_config",
@@ -516,7 +520,7 @@ mod tests {
 
     #[test]
     fn test_unstable_ops_count() {
-        assert_eq!(UNSTABLE_OPS.len(), 186);
+        assert_eq!(UNSTABLE_OPS.len(), 190);
     }
 
     #[test]

src/commands/llm_obs.rs

@@ -3,10 +3,11 @@ use datadog_api_client::datadogV2::api_llm_observability::{
     LLMObservabilityAPI, ListLLMObsAnnotationQueuesOptionalParams,
 };
 use datadog_api_client::datadogV2::model::{
-    LLMObsAnnotationQueueInteractionsRequest, LLMObsAnnotationQueueRequest,
-    LLMObsAnnotationQueueUpdateRequest, LLMObsCustomEvalConfigUpdateRequest,
-    LLMObsDatasetBatchUpdateRequest, LLMObsDatasetCloneRequest, LLMObsDatasetRequest,
-    LLMObsDatasetRestoreVersionRequest, LLMObsDeleteAnnotationQueueInteractionsRequest,
+    LLMObsAnnotationQueueInteractionsRequest, LLMObsAnnotationQueueLabelSchemaUpdateRequest,
+    LLMObsAnnotationQueueRequest, LLMObsAnnotationQueueUpdateRequest, LLMObsAnnotationsRequest,
+    LLMObsCustomEvalConfigUpdateRequest, LLMObsDatasetBatchUpdateRequest,
+    LLMObsDatasetCloneRequest, LLMObsDatasetRequest, LLMObsDatasetRestoreVersionRequest,
+    LLMObsDeleteAnnotationQueueInteractionsRequest, LLMObsDeleteAnnotationsRequest,
     LLMObsDeleteExperimentsRequest, LLMObsProjectRequest,
 };
 
@@ -596,6 +597,57 @@ pub async fn annotation_queue_interactions_list(cfg: &Config, queue_id: &str) ->
     formatter::output(cfg, &resp)
 }
 
+pub async fn annotation_queue_schema_get(cfg: &Config, queue_id: &str) -> Result<()> {
+    let api = make_api(cfg);
+    let resp = api
+        .get_llm_obs_annotation_queue_label_schema(queue_id.to_string())
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to get annotation queue label schema: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+pub async fn annotation_queue_schema_update(
+    cfg: &Config,
+    queue_id: &str,
+    file: &str,
+) -> Result<()> {
+    let body: LLMObsAnnotationQueueLabelSchemaUpdateRequest = util::read_json_file(file)?;
+    let api = make_api(cfg);
+    let resp = api
+        .update_llm_obs_annotation_queue_label_schema(queue_id.to_string(), body)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to update annotation queue label schema: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+pub async fn annotation_queue_annotations_upsert(
+    cfg: &Config,
+    queue_id: &str,
+    file: &str,
+) -> Result<()> {
+    let body: LLMObsAnnotationsRequest = util::read_json_file(file)?;
+    let api = make_api(cfg);
+    let resp = api
+        .upsert_llm_obs_annotations(queue_id.to_string(), body)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to upsert annotations: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+pub async fn annotation_queue_annotations_delete(
+    cfg: &Config,
+    queue_id: &str,
+    file: &str,
+) -> Result<()> {
+    let body: LLMObsDeleteAnnotationsRequest = util::read_json_file(file)?;
+    let api = make_api(cfg);
+    let resp = api
+        .delete_llm_obs_annotations(queue_id.to_string(), body)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to delete annotations: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
 // ---- Custom Evaluator Configs ----
 
 pub async fn eval_config_get(cfg: &Config, eval_name: &str) -> Result<()> {
@@ -4961,4 +5013,173 @@ mod tests {
         assert!(result.is_err(), "should fail on 500");
         cleanup_env();
     }
+
+    // ---- Annotation queue label schemas ----
+
+    const LABEL_SCHEMA_BODY: &str = r#"{"data":{"id":"queue-1","type":"queues","attributes":{"annotation_schema":{"label_schemas":[{"id":"ls-1","name":"quality","type":"score","min":0.0,"max":5.0,"is_required":true}]}}}}"#;
+
+    #[tokio::test]
+    async fn test_annotation_queue_schema_get() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+        let _mock = mock_any(&mut server, "GET", LABEL_SCHEMA_BODY).await;
+
+        let result = super::annotation_queue_schema_get(&cfg, "queue-1").await;
+        assert!(result.is_ok(), "schema_get failed: {:?}", result.err());
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_annotation_queue_schema_get_404() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+        let _mock = server
+            .mock("GET", mockito::Matcher::Any)
+            .match_query(mockito::Matcher::Any)
+            .with_status(404)
+            .with_header("content-type", "application/json")
+            .with_body(r#"{"errors":["queue not found"]}"#)
+            .create_async()
+            .await;
+
+        let result = super::annotation_queue_schema_get(&cfg, "missing-queue").await;
+        assert!(result.is_err(), "should fail on 404");
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_annotation_queue_schema_update() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+        let _mock = mock_any(&mut server, "PUT", LABEL_SCHEMA_BODY).await;
+
+        let path = write_temp_json(
+            "pup_aq_schema_update.json",
+            r#"{"data":{"type":"queues","attributes":{"annotation_schema":{"label_schemas":[{"name":"quality","type":"score","min":0.0,"max":5.0}]}}}}"#,
+        );
+        let result =
+            super::annotation_queue_schema_update(&cfg, "queue-1", path.to_str().unwrap()).await;
+        assert!(result.is_ok(), "schema_update failed: {:?}", result.err());
+        let _ = std::fs::remove_file(&path);
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_annotation_queue_schema_update_missing_file() {
+        let _lock = lock_env().await;
+        let server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let result =
+            super::annotation_queue_schema_update(&cfg, "queue-1", "/nonexistent/schema.json")
+                .await;
+        let err = result
+            .expect_err("should fail on unreadable file")
+            .to_string();
+        assert!(err.contains("failed to read file"), "unexpected: {err}");
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_annotation_queue_schema_update_malformed_json() {
+        let _lock = lock_env().await;
+        let server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        // Valid JSON but missing the required `data` member.
+        let path = write_temp_json("pup_aq_schema_bad.json", r#"{"nope":true}"#);
+        let result =
+            super::annotation_queue_schema_update(&cfg, "queue-1", path.to_str().unwrap()).await;
+        let err = result.expect_err("should reject bad body").to_string();
+        assert!(err.contains("failed to parse JSON"), "unexpected: {err}");
+        let _ = std::fs::remove_file(&path);
+        cleanup_env();
+    }
+
+    // ---- Annotations on queue interactions ----
+
+    #[tokio::test]
+    async fn test_annotation_queue_annotations_upsert() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+        let body = r#"{"data":{"id":"upsert-1","type":"annotations","attributes":{"annotations":[{"id":"a-1","interaction_id":"i-1","created_at":"2024-01-01T00:00:00Z","created_by":"user-1","modified_at":"2024-01-01T00:00:00Z","modified_by":"user-1","label_values":[{"label_schema_id":"ls-1","value":4.0}]}]}}}"#;
+        let _mock = mock_any(&mut server, "POST", body).await;
+
+        let path = write_temp_json(
+            "pup_aq_annotations_upsert.json",
+            r#"{"data":{"type":"annotations","attributes":{"annotations":[{"interaction_id":"i-1","label_values":[{"label_schema_id":"ls-1","value":4.0}]}]}}}"#,
+        );
+        let result =
+            super::annotation_queue_annotations_upsert(&cfg, "queue-1", path.to_str().unwrap())
+                .await;
+        assert!(result.is_ok(), "upsert failed: {:?}", result.err());
+        let _ = std::fs::remove_file(&path);
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_annotation_queue_annotations_upsert_400() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+        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":["unknown label_schema_id"]}"#)
+            .create_async()
+            .await;
+
+        let path = write_temp_json(
+            "pup_aq_annotations_upsert_400.json",
+            r#"{"data":{"type":"annotations","attributes":{"annotations":[{"interaction_id":"i-1","label_values":[{"label_schema_id":"bogus","value":4.0}]}]}}}"#,
+        );
+        let result =
+            super::annotation_queue_annotations_upsert(&cfg, "queue-1", path.to_str().unwrap())
+                .await;
+        assert!(result.is_err(), "should fail on 400");
+        let _ = std::fs::remove_file(&path);
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_annotation_queue_annotations_delete() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+        let body = r#"{"data":{"id":"delete-1","type":"annotations","attributes":{"annotation_ids":["a-1"],"errors":[]}}}"#;
+        let _mock = mock_any(&mut server, "POST", body).await;
+
+        let path = write_temp_json(
+            "pup_aq_annotations_delete.json",
+            r#"{"data":{"type":"annotations","attributes":{"annotation_ids":["a-1"]}}}"#,
+        );
+        let result =
+            super::annotation_queue_annotations_delete(&cfg, "queue-1", path.to_str().unwrap())
+                .await;
+        assert!(result.is_ok(), "delete failed: {:?}", result.err());
+        let _ = std::fs::remove_file(&path);
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_annotation_queue_annotations_delete_malformed_json() {
+        let _lock = lock_env().await;
+        let server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let path = write_temp_json("pup_aq_annotations_delete_bad.json", r#"{ not json"#);
+        let result =
+            super::annotation_queue_annotations_delete(&cfg, "queue-1", path.to_str().unwrap())
+                .await;
+        let err = result.expect_err("should reject bad JSON").to_string();
+        assert!(err.contains("failed to parse JSON"), "unexpected: {err}");
+        let _ = std::fs::remove_file(&path);
+        cleanup_env();
+    }
 }

src/main.rs

@@ -9886,6 +9886,50 @@ enum LlmObsAnnotationQueuesActions {
         #[command(subcommand)]
         action: LlmObsAnnotationQueueInteractionsActions,
     },
+    /// Manage the label schema of an annotation queue
+    Schema {
+        #[command(subcommand)]
+        action: LlmObsAnnotationQueueSchemaActions,
+    },
+    /// Manage annotations on interactions in an annotation queue
+    Annotations {
+        #[command(subcommand)]
+        action: LlmObsAnnotationQueueAnnotationsActions,
+    },
+}
+
+#[derive(Subcommand)]
+enum LlmObsAnnotationQueueSchemaActions {
+    /// Get the label schema for an annotation queue
+    Get {
+        #[arg(help = "Annotation queue ID")]
+        queue_id: String,
+    },
+    /// Replace the label schema for an annotation queue
+    Update {
+        #[arg(help = "Annotation queue ID")]
+        queue_id: String,
+        #[arg(long, help = "JSON file with label schema update body (required)")]
+        file: String,
+    },
+}
+
+#[derive(Subcommand)]
+enum LlmObsAnnotationQueueAnnotationsActions {
+    /// Create or update annotations on interactions in an annotation queue
+    Upsert {
+        #[arg(help = "Annotation queue ID")]
+        queue_id: String,
+        #[arg(long, help = "JSON file with annotations body (required)")]
+        file: String,
+    },
+    /// Delete annotations from interactions in an annotation queue
+    Delete {
+        #[arg(help = "Annotation queue ID")]
+        queue_id: String,
+        #[arg(long, help = "JSON file with annotations to delete (required)")]
+        file: String,
+    },
 }
 
 #[derive(Subcommand)]
@@ -17032,6 +17076,31 @@ async fn main_inner() -> anyhow::Result<()> {
                                 .await?;
                         }
                     },
+                    LlmObsAnnotationQueuesActions::Schema { action } => match action {
+                        LlmObsAnnotationQueueSchemaActions::Get { queue_id } => {
+                            commands::llm_obs::annotation_queue_schema_get(&cfg, &queue_id).await?;
+                        }
+                        LlmObsAnnotationQueueSchemaActions::Update { queue_id, file } => {
+                            commands::llm_obs::annotation_queue_schema_update(
+                                &cfg, &queue_id, &file,
+                            )
+                            .await?;
+                        }
+                    },
+                    LlmObsAnnotationQueuesActions::Annotations { action } => match action {
+                        LlmObsAnnotationQueueAnnotationsActions::Upsert { queue_id, file } => {
+                            commands::llm_obs::annotation_queue_annotations_upsert(
+                                &cfg, &queue_id, &file,
+                            )
+                            .await?;
+                        }
+                        LlmObsAnnotationQueueAnnotationsActions::Delete { queue_id, file } => {
+                            commands::llm_obs::annotation_queue_annotations_delete(
+                                &cfg, &queue_id, &file,
+                            )
+                            .await?;
+                        }
+                    },
                 },
                 LlmObsActions::EvalConfig { action } => match action {
                     LlmObsEvalConfigActions::Get { eval_name } => {