diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index 1ddab045..ee1623d6 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, diff, 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), annotation-queues (create, list, update, delete, interactions (add, delete, list), schema (get, update), annotations (upsert, delete)), 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), annotations (export), 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 | ✅ | @@ -94,6 +94,18 @@ pup [options] # Nested commands ## Common Patterns +### Export an LLM Observability annotated interaction + +Export one annotation together with all event data for its trace, span, experiment trace, or session. `--queue` accepts either the queue ID or its exact name. Without `--out`, the command uses Pup's standard output formatter (JSON by default) and honors global `--output` and `--jq` flags. File exports default to JSON; pass `--format jsonl` for JSONL. Existing files are not replaced unless `--force` is provided. + +```bash +pup llm-obs annotations export \ + --queue quality-review \ + --interaction-id 23851556-a8c7-41c1-be75-03eb665a132f \ + --format jsonl \ + --out interaction.jsonl +``` + ### List Operations ```bash pup list [--flags] diff --git a/docs/EXAMPLES.md b/docs/EXAMPLES.md index 3fd5e817..e5529629 100644 --- a/docs/EXAMPLES.md +++ b/docs/EXAMPLES.md @@ -26,6 +26,28 @@ export DD_APP_KEY="your-app-key" export DD_SITE="datadoghq.com" ``` +## LLM Observability + +### Export an Annotated Interaction +```bash +# Print the annotation and complete interaction data as JSON +pup llm-obs annotations export \ + --queue quality-review \ + --interaction-id 23851556-a8c7-41c1-be75-03eb665a132f + +# Select part of the standard JSON output +pup --jq '.interaction_data.events' llm-obs annotations export \ + --queue 13851556-a8c7-41c1-be75-03eb665a132f \ + --interaction-id 23851556-a8c7-41c1-be75-03eb665a132f + +# Write JSONL to a file; --force is required to replace an existing file +pup llm-obs annotations export \ + --queue quality-review \ + --interaction-id 23851556-a8c7-41c1-be75-03eb665a132f \ + --out interaction.jsonl \ + --format jsonl +``` + ## Metrics ### List Metrics diff --git a/src/commands/llm_obs.rs b/src/commands/llm_obs.rs index dca0652b..317109f6 100644 --- a/src/commands/llm_obs.rs +++ b/src/commands/llm_obs.rs @@ -16,6 +16,11 @@ use crate::raw_client; use crate::util; use crate::util_ext; +use std::collections::HashSet; +use std::fs::OpenOptions; +use std::io::Write; +use std::path::{Path, PathBuf}; + fn make_api(cfg: &Config) -> AgentObservabilityAPI { crate::make_api!(AgentObservabilityAPI, cfg) } @@ -596,6 +601,267 @@ pub async fn annotation_queue_interactions_list(cfg: &Config, queue_id: &str) -> formatter::output(cfg, &resp) } +const ANNOTATED_INTERACTION_DATA_PAGE_SIZE: &str = "500"; + +#[derive(Debug, serde::Deserialize)] +struct AnnotatedInteractionDataPage { + annotated_interaction: serde_json::Value, + interaction_data: InteractionDataPage, +} + +#[derive(Debug, serde::Deserialize)] +struct InteractionDataPage { + #[serde(rename = "type")] + interaction_type: String, + events: Vec, + #[serde(default)] + next_cursor: Option, +} + +fn parse_annotated_interaction_data_page( + response: serde_json::Value, +) -> Result { + let attributes = response + .get("data") + .and_then(|data| data.get("attributes")) + .cloned() + .ok_or_else(|| anyhow::anyhow!("response is missing data.attributes"))?; + serde_json::from_value(attributes) + .map_err(|e| anyhow::anyhow!("invalid annotated interaction response: {e}")) +} + +fn select_annotation_queue_id(response: &serde_json::Value, queue: &str) -> Result { + let queues = response + .get("data") + .and_then(serde_json::Value::as_array) + .ok_or_else(|| anyhow::anyhow!("annotation queue list response is missing data"))?; + + let matching_ids: Vec<&str> = queues + .iter() + .filter(|item| { + item.pointer("/attributes/name") + .and_then(serde_json::Value::as_str) + == Some(queue) + }) + .filter_map(|item| { + item.pointer("/attributes/queue_id") + .and_then(serde_json::Value::as_str) + .or_else(|| item.get("id").and_then(serde_json::Value::as_str)) + }) + .collect(); + + match matching_ids.as_slice() { + [queue_id] if uuid::Uuid::parse_str(queue_id).is_ok() => Ok((*queue_id).to_string()), + [_] => anyhow::bail!("annotation queue '{queue}' returned an invalid queue ID"), + [] => anyhow::bail!("no annotation queue found with exact name '{queue}'"), + _ => anyhow::bail!( + "multiple annotation queues are named '{queue}'; pass the queue ID instead" + ), + } +} + +async fn resolve_annotation_queue_id(cfg: &Config, queue: &str) -> Result { + if uuid::Uuid::parse_str(queue).is_ok() { + return Ok(queue.to_string()); + } + if queue.is_empty() { + anyhow::bail!("--queue cannot be empty"); + } + + let response = raw_client::raw_get(cfg, "/api/v2/llm-obs/v1/annotation-queues", &[]) + .await + .map_err(|e| { + anyhow::anyhow!( + "failed to resolve annotation queue '{queue}'; ensure the caller has the llm_observability_read scope and access to the queue: {e:?}" + ) + })?; + select_annotation_queue_id(&response, queue) +} + +fn merge_annotated_interaction_page( + annotated_interaction: &mut Option, + interaction_type: &mut Option, + events: &mut Vec, + page: AnnotatedInteractionDataPage, +) -> Result> { + if let Some(existing) = annotated_interaction.as_ref() { + if existing != &page.annotated_interaction { + anyhow::bail!("annotated interaction changed while the export was being fetched"); + } + } else { + *annotated_interaction = Some(page.annotated_interaction); + } + if let Some(existing) = interaction_type.as_ref() { + if existing != &page.interaction_data.interaction_type { + anyhow::bail!("interaction type changed while the export was being fetched"); + } + } else { + *interaction_type = Some(page.interaction_data.interaction_type); + } + + events.extend(page.interaction_data.events); + Ok(page + .interaction_data + .next_cursor + .filter(|value| !value.is_empty())) +} + +async fn fetch_annotated_interaction_export( + cfg: &Config, + queue_id: &str, + interaction_id: &str, +) -> Result { + let path = format!( + "/api/v2/llm-obs/v1/annotation-queues/{queue_id}/annotated-interactions/{interaction_id}" + ); + let mut cursor: Option = None; + let mut seen_cursors = HashSet::new(); + let mut annotated_interaction: Option = None; + let mut interaction_type: Option = None; + let mut events = Vec::new(); + + loop { + let mut query = vec![("limit", ANNOTATED_INTERACTION_DATA_PAGE_SIZE)]; + if let Some(ref value) = cursor { + query.push(("cursor", value.as_str())); + } + let response = raw_client::raw_get(cfg, &path, &query) + .await + .map_err(|e| { + anyhow::anyhow!( + "failed to fetch annotated interaction data; ensure the caller has the llm_observability_read scope and access to the queue: {e:?}" + ) + })?; + let page = parse_annotated_interaction_data_page(response)?; + let Some(next_cursor) = merge_annotated_interaction_page( + &mut annotated_interaction, + &mut interaction_type, + &mut events, + page, + )? + else { + break; + }; + if !seen_cursors.insert(next_cursor.clone()) { + anyhow::bail!("the server returned a repeated pagination cursor"); + } + cursor = Some(next_cursor); + } + + Ok(serde_json::json!({ + "annotated_interaction": annotated_interaction + .ok_or_else(|| anyhow::anyhow!("response did not contain an annotated interaction"))?, + "interaction_data": { + "type": interaction_type + .ok_or_else(|| anyhow::anyhow!("response did not contain an interaction type"))?, + "events": events, + }, + })) +} + +fn output_annotation_export( + cfg: &Config, + export: &serde_json::Value, + format: Option<&str>, + out: Option<&str>, + force: bool, +) -> Result<()> { + if let Some(path) = out { + let bytes = serialize_annotation_export(export, format.unwrap_or("json"))?; + write_annotation_export(Path::new(path), &bytes, force)?; + eprintln!("Exported annotated interaction to {path}"); + return Ok(()); + } + formatter::output(cfg, export) +} + +fn serialize_annotation_export(value: &serde_json::Value, format: &str) -> Result> { + let mut bytes = match format { + "json" => serde_json::to_vec_pretty(value)?, + "jsonl" => serde_json::to_vec(value)?, + _ => anyhow::bail!("unsupported export format '{format}'"), + }; + bytes.push(b'\n'); + Ok(bytes) +} + +fn temporary_export_path(path: &Path, attempt: u8) -> Result { + let name = path + .file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| anyhow::anyhow!("output path must name a file"))?; + Ok(path.with_file_name(format!(".{name}.pup-{}-{attempt}.tmp", std::process::id()))) +} + +fn write_annotation_export(path: &Path, bytes: &[u8], force: bool) -> Result<()> { + if path.exists() && !force { + anyhow::bail!( + "output file '{}' already exists; pass --force to overwrite it", + path.display() + ); + } + + let mut temporary = None; + for attempt in 0..100 { + let candidate = temporary_export_path(path, attempt)?; + let mut options = OpenOptions::new(); + options.write(true).create_new(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); + } + match options.open(&candidate) { + Ok(file) => { + temporary = Some((candidate, file)); + break; + } + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue, + Err(error) => { + return Err(anyhow::anyhow!( + "failed to create export beside '{}': {error}", + path.display() + )) + } + } + } + let (temporary_path, mut file) = + temporary.ok_or_else(|| anyhow::anyhow!("could not allocate a temporary export file"))?; + + let result = (|| -> Result<()> { + file.write_all(bytes)?; + file.sync_all()?; + drop(file); + #[cfg(windows)] + if force && path.exists() { + std::fs::remove_file(path)?; + } + std::fs::rename(&temporary_path, path)?; + Ok(()) + })(); + if result.is_err() { + let _ = std::fs::remove_file(&temporary_path); + } + result.map_err(|e| anyhow::anyhow!("failed to write export '{}': {e}", path.display())) +} + +pub async fn annotations_export( + cfg: &Config, + queue: &str, + interaction_id: &str, + format: Option<&str>, + out: Option<&str>, + force: bool, +) -> Result<()> { + if uuid::Uuid::parse_str(interaction_id).is_err() { + anyhow::bail!("--interaction-id must be a UUID"); + } + let queue_id = resolve_annotation_queue_id(cfg, queue).await?; + let export = fetch_annotated_interaction_export(cfg, &queue_id, interaction_id).await?; + + output_annotation_export(cfg, &export, format, out, force) +} + /// Uses the raw client rather than the typed one: a queue with no schema yet answers with /// `"annotation_schema": null`, but the generated client models it as a non-`Option` /// `LLMObsAnnotationSchema` and fails with "invalid type: null, expected a mapping". Most queues @@ -1446,6 +1712,405 @@ mod tests { std::env::remove_var("DD_TOKEN_STORAGE"); } + #[tokio::test] + async fn test_annotations_export_fetches_all_pages() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + let queue_id = "13851556-a8c7-41c1-be75-03eb665a132f"; + let interaction_id = "23851556-a8c7-41c1-be75-03eb665a132f"; + let path = format!( + "/api/v2/llm-obs/v1/annotation-queues/{queue_id}/annotated-interactions/{interaction_id}" + ); + let annotated_interaction = serde_json::json!({ + "interaction_id": interaction_id, + "queue_id": queue_id, + "queue_name": "quality-review", + "can_annotate": true, + "annotations": [{"label": "quality", "value": "good"}], + }); + let first_body = serde_json::json!({ + "data": { + "id": interaction_id, + "type": "annotated_interaction_data", + "attributes": { + "annotated_interaction": annotated_interaction, + "interaction_data": { + "type": "session", + "events": [{"content": {"span_id": "span-1"}}], + "next_cursor": "cursor-1", + } + } + } + }); + let second_body = serde_json::json!({ + "data": { + "id": interaction_id, + "type": "annotated_interaction_data", + "attributes": { + "annotated_interaction": annotated_interaction, + "interaction_data": { + "type": "session", + "events": [{"content": {"span_id": "span-2", "expected_output": "ok"}}] + } + } + } + }); + let _first = server + .mock("GET", path.as_str()) + .match_query(mockito::Matcher::Exact("limit=500".into())) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(first_body.to_string()) + .create_async() + .await; + let _second = server + .mock("GET", path.as_str()) + .match_query(mockito::Matcher::Exact("limit=500&cursor=cursor-1".into())) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(second_body.to_string()) + .create_async() + .await; + + let export = super::fetch_annotated_interaction_export(&cfg, queue_id, interaction_id) + .await + .expect("multi-page export should succeed"); + + assert_eq!(export["annotated_interaction"], annotated_interaction); + assert_eq!(export["interaction_data"]["type"], "session"); + assert_eq!( + export["interaction_data"]["events"] + .as_array() + .unwrap() + .len(), + 2 + ); + assert_eq!( + export["interaction_data"]["events"][1]["content"]["expected_output"], + "ok" + ); + cleanup_env(); + } + + #[tokio::test] + async fn test_annotations_export_resolves_exact_queue_name() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + let queue_id = "13851556-a8c7-41c1-be75-03eb665a132f"; + let interaction_id = "23851556-a8c7-41c1-be75-03eb665a132f"; + let _queues = server + .mock("GET", "/api/v2/llm-obs/v1/annotation-queues") + .match_query(mockito::Matcher::Missing) + .with_status(200) + .with_header("content-type", "application/json") + .with_body( + serde_json::json!({ + "data": [{ + "id": queue_id, + "type": "queues", + "attributes": {"queue_id": queue_id, "name": "quality-review"} + }] + }) + .to_string(), + ) + .create_async() + .await; + let interaction_path = format!( + "/api/v2/llm-obs/v1/annotation-queues/{queue_id}/annotated-interactions/{interaction_id}" + ); + let _interaction = server + .mock("GET", interaction_path.as_str()) + .match_query(mockito::Matcher::Exact("limit=500".into())) + .with_status(200) + .with_header("content-type", "application/json") + .with_body( + serde_json::json!({ + "data": {"attributes": { + "annotated_interaction": {"interaction_id": interaction_id}, + "interaction_data": {"type": "experiment_trace", "events": []} + }} + }) + .to_string(), + ) + .create_async() + .await; + let temp = TempDir::new("annotation_export"); + let output = temp.path().join("interaction.jsonl"); + + super::annotations_export( + &cfg, + "quality-review", + interaction_id, + Some("jsonl"), + output.to_str(), + false, + ) + .await + .expect("named queue export should succeed"); + + let exported: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(output).expect("export should be readable"), + ) + .expect("export should contain JSON"); + assert_eq!(exported["interaction_data"]["type"], "experiment_trace"); + cleanup_env(); + } + + #[tokio::test] + async fn test_annotations_export_rejects_repeated_cursor() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + let queue_id = "13851556-a8c7-41c1-be75-03eb665a132f"; + let interaction_id = "23851556-a8c7-41c1-be75-03eb665a132f"; + let path = format!( + "/api/v2/llm-obs/v1/annotation-queues/{queue_id}/annotated-interactions/{interaction_id}" + ); + let body = serde_json::json!({ + "data": {"attributes": { + "annotated_interaction": {"interaction_id": interaction_id}, + "interaction_data": { + "type": "trace", + "events": [], + "next_cursor": "same-cursor" + } + }} + }); + let _first = server + .mock("GET", path.as_str()) + .match_query(mockito::Matcher::Exact("limit=500".into())) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body.to_string()) + .create_async() + .await; + let _second = server + .mock("GET", path.as_str()) + .match_query(mockito::Matcher::Exact( + "limit=500&cursor=same-cursor".into(), + )) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body.to_string()) + .create_async() + .await; + + let error = super::fetch_annotated_interaction_export(&cfg, queue_id, interaction_id) + .await + .expect_err("repeated cursor should fail"); + assert!(error.to_string().contains("repeated pagination cursor")); + cleanup_env(); + } + + #[test] + fn test_annotation_export_refuses_overwrite_without_force() { + let temp = TempDir::new("annotation_export_overwrite"); + let output = temp.path().join("interaction.jsonl"); + std::fs::write(&output, b"existing\n").unwrap(); + + let error = super::write_annotation_export(&output, b"replacement\n", false) + .expect_err("existing output should be protected"); + + assert!(error.to_string().contains("--force")); + assert_eq!(std::fs::read(&output).unwrap(), b"existing\n"); + super::write_annotation_export(&output, b"replacement\n", true) + .expect("--force should replace the output"); + assert_eq!(std::fs::read(&output).unwrap(), b"replacement\n"); + } + + #[test] + fn test_parse_annotated_interaction_data_page_rejects_missing_attributes() { + let error = super::parse_annotated_interaction_data_page(serde_json::json!({"data": {}})) + .expect_err("missing attributes should fail"); + assert!(error.to_string().contains("data.attributes")); + } + + #[test] + fn test_select_annotation_queue_id_rejects_invalid_queue_responses() { + let missing_data = super::select_annotation_queue_id(&serde_json::json!({}), "review") + .expect_err("missing data should fail"); + assert!(missing_data.to_string().contains("missing data")); + + let not_found = + super::select_annotation_queue_id(&serde_json::json!({"data": []}), "review") + .expect_err("unknown queue should fail"); + assert!(not_found.to_string().contains("no annotation queue")); + + let invalid_id = super::select_annotation_queue_id( + &serde_json::json!({"data": [{ + "id": "not-a-uuid", + "attributes": {"name": "review"} + }]}), + "review", + ) + .expect_err("invalid queue ID should fail"); + assert!(invalid_id.to_string().contains("invalid queue ID")); + + let duplicate = super::select_annotation_queue_id( + &serde_json::json!({"data": [ + {"id": "13851556-a8c7-41c1-be75-03eb665a132f", "attributes": {"name": "review"}}, + {"id": "23851556-a8c7-41c1-be75-03eb665a132f", "attributes": {"name": "review"}} + ]}), + "review", + ) + .expect_err("duplicate queue names should fail"); + assert!(duplicate.to_string().contains("multiple annotation queues")); + } + + #[tokio::test] + async fn test_resolve_annotation_queue_id_rejects_empty_name() { + let cfg = test_config("http://unused.local"); + let error = super::resolve_annotation_queue_id(&cfg, "") + .await + .expect_err("empty queue name should fail before a request"); + assert!(error.to_string().contains("--queue cannot be empty")); + } + + #[test] + fn test_merge_annotated_interaction_page_rejects_changes_between_pages() { + let mut annotation = Some(serde_json::json!({"interaction_id": "one"})); + let mut interaction_type = Some("trace".to_string()); + let mut events = Vec::new(); + let changed_annotation = super::AnnotatedInteractionDataPage { + annotated_interaction: serde_json::json!({"interaction_id": "two"}), + interaction_data: super::InteractionDataPage { + interaction_type: "trace".to_string(), + events: vec![], + next_cursor: None, + }, + }; + let error = super::merge_annotated_interaction_page( + &mut annotation, + &mut interaction_type, + &mut events, + changed_annotation, + ) + .expect_err("changed annotation should fail"); + assert!(error.to_string().contains("annotated interaction changed")); + + let changed_type = super::AnnotatedInteractionDataPage { + annotated_interaction: serde_json::json!({"interaction_id": "one"}), + interaction_data: super::InteractionDataPage { + interaction_type: "session".to_string(), + events: vec![], + next_cursor: None, + }, + }; + let error = super::merge_annotated_interaction_page( + &mut annotation, + &mut interaction_type, + &mut events, + changed_type, + ) + .expect_err("changed interaction type should fail"); + assert!(error.to_string().contains("interaction type changed")); + } + + #[test] + fn test_annotation_export_output_honors_standard_jq_filter() { + let mut cfg = test_config("http://unused.local"); + cfg.jq = Some("[".to_string()); + let error = super::output_annotation_export( + &cfg, + &serde_json::json!({"interaction_data": {"events": []}}), + None, + None, + false, + ) + .expect_err("invalid global jq filter should be applied and rejected"); + assert!(error.to_string().contains("jq")); + } + + #[test] + fn test_annotation_export_rejects_unsupported_file_format() { + let error = super::serialize_annotation_export(&serde_json::json!({}), "yaml") + .expect_err("unsupported file format should fail"); + assert!(error.to_string().contains("unsupported export format")); + } + + #[test] + fn test_annotation_export_rejects_path_without_file_name() { + let error = super::temporary_export_path(std::path::Path::new("/"), 0) + .expect_err("directory path should fail"); + assert!(error.to_string().contains("must name a file")); + } + + #[test] + fn test_annotation_export_file_defaults_to_json() { + let cfg = test_config("http://unused.local"); + let temp = TempDir::new("annotation_export_default_json"); + let output = temp.path().join("interaction.json"); + super::output_annotation_export( + &cfg, + &serde_json::json!({"interaction_data": {"events": []}}), + None, + output.to_str(), + false, + ) + .expect("file export should default to JSON"); + let contents = std::fs::read_to_string(output).expect("export should be readable"); + assert!( + contents.starts_with("{\n"), + "JSON export should be pretty-printed" + ); + } + + #[test] + fn test_annotation_export_reports_file_creation_failure() { + let temp = TempDir::new("annotation_export_missing_parent"); + let output = temp.path().join("missing").join("interaction.json"); + let error = super::write_annotation_export(&output, b"{}\n", false) + .expect_err("missing parent directory should fail"); + assert!(error.to_string().contains("failed to create export")); + } + + #[tokio::test] + async fn test_annotation_export_api_error_names_required_scope() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + let queue_id = "13851556-a8c7-41c1-be75-03eb665a132f"; + let interaction_id = "23851556-a8c7-41c1-be75-03eb665a132f"; + let path = format!( + "/api/v2/llm-obs/v1/annotation-queues/{queue_id}/annotated-interactions/{interaction_id}" + ); + let _request = server + .mock("GET", path.as_str()) + .match_query(mockito::Matcher::Exact("limit=500".into())) + .with_status(403) + .with_header("content-type", "application/json") + .with_body(r#"{"errors":["forbidden"]}"#) + .create_async() + .await; + + let error = super::fetch_annotated_interaction_export(&cfg, queue_id, interaction_id) + .await + .expect_err("API rejection should include remediation"); + assert!(error.to_string().contains("llm_observability_read")); + assert!(error.to_string().contains("access to the queue")); + cleanup_env(); + } + + #[tokio::test] + async fn test_annotations_export_rejects_invalid_interaction_id() { + let cfg = test_config("http://unused.local"); + let error = super::annotations_export( + &cfg, + "13851556-a8c7-41c1-be75-03eb665a132f", + "not-a-uuid", + None, + None, + false, + ) + .await + .expect_err("invalid interaction ID should fail before a request"); + assert!(error + .to_string() + .contains("--interaction-id must be a UUID")); + } + #[tokio::test] async fn test_llm_obs_projects_list_with_filters() { let _lock = lock_env().await; diff --git a/src/main.rs b/src/main.rs index 7ac8e9f6..972fbc08 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9487,6 +9487,11 @@ enum LlmObsActions { #[command(subcommand)] action: LlmObsSpansActions, }, + /// Export an annotated interaction with its complete interaction data + Annotations { + #[command(subcommand)] + action: LlmObsAnnotationsActions, + }, /// Manage LLM Observability annotation queues #[command(name = "annotation-queues")] AnnotationQueues { @@ -9529,6 +9534,28 @@ enum LlmObsActions { }, } +#[derive(Subcommand)] +enum LlmObsAnnotationsActions { + /// Export one annotated interaction and all of its event data + Export { + #[arg(long, help = "Annotation queue ID or exact queue name")] + queue: String, + #[arg(long, help = "Interaction ID")] + interaction_id: String, + #[arg( + long, + value_parser = ["json", "jsonl"], + requires = "out", + help = "Output-file format: json (default) or jsonl; requires --out" + )] + format: Option, + #[arg(long, help = "Output file; writes to stdout when omitted")] + out: Option, + #[arg(long, requires = "out", help = "Overwrite an existing output file")] + force: bool, + }, +} + #[derive(Subcommand)] enum LlmObsAgentInsightsActions { /// List Agent Insights as compact triage cards @@ -17509,6 +17536,25 @@ async fn main_inner() -> anyhow::Result<()> { } }, }, + LlmObsActions::Annotations { action } => match action { + LlmObsAnnotationsActions::Export { + queue, + interaction_id, + format, + out, + force, + } => { + commands::llm_obs::annotations_export( + &cfg, + &queue, + &interaction_id, + format.as_deref(), + out.as_deref(), + force, + ) + .await?; + } + }, LlmObsActions::EvalConfig { action } => match action { LlmObsEvalConfigActions::Get { eval_name } => { commands::llm_obs::eval_config_get(&cfg, &eval_name).await?; diff --git a/src/test_commands.rs b/src/test_commands.rs index f9858eb8..af1f4805 100644 --- a/src/test_commands.rs +++ b/src/test_commands.rs @@ -8,6 +8,104 @@ use clap::CommandFactory; +#[test] +fn test_llm_obs_annotations_export_parses() { + use clap::Parser; + + let cli = crate::Cli::try_parse_from([ + "pup", + "llm-obs", + "annotations", + "export", + "--queue", + "quality-review", + "--interaction-id", + "23851556-a8c7-41c1-be75-03eb665a132f", + "--format", + "jsonl", + "--out", + "interaction.jsonl", + "--force", + ]) + .expect("LLM Observability annotation export should parse"); + + let crate::Commands::LlmObs { action } = cli.command else { + panic!("expected Commands::LlmObs"); + }; + let crate::LlmObsActions::Annotations { action } = action else { + panic!("expected LlmObsActions::Annotations"); + }; + let crate::LlmObsAnnotationsActions::Export { + queue, + interaction_id, + format, + out, + force, + } = action; + assert_eq!(queue, "quality-review"); + assert_eq!(interaction_id, "23851556-a8c7-41c1-be75-03eb665a132f"); + assert_eq!(format.as_deref(), Some("jsonl")); + assert_eq!(out.as_deref(), Some("interaction.jsonl")); + assert!(force); +} + +#[test] +fn test_llm_obs_annotations_export_defaults_to_standard_json_output() { + use clap::Parser; + + let cli = crate::Cli::try_parse_from([ + "pup", + "llm-obs", + "annotations", + "export", + "--queue", + "quality-review", + "--interaction-id", + "23851556-a8c7-41c1-be75-03eb665a132f", + ]) + .expect("annotation export should parse without file flags"); + + let crate::Commands::LlmObs { action } = cli.command else { + panic!("expected Commands::LlmObs"); + }; + let crate::LlmObsActions::Annotations { action } = action else { + panic!("expected LlmObsActions::Annotations"); + }; + let crate::LlmObsAnnotationsActions::Export { + format, out, force, .. + } = action; + assert!(format.is_none()); + assert!(out.is_none()); + assert!(!force); +} + +#[test] +fn test_llm_obs_annotations_export_file_flags_require_out() { + use clap::Parser; + + for flag in ["--format", "--force"] { + let mut args = vec![ + "pup", + "llm-obs", + "annotations", + "export", + "--queue", + "quality-review", + "--interaction-id", + "23851556-a8c7-41c1-be75-03eb665a132f", + flag, + ]; + if flag == "--format" { + args.push("jsonl"); + } + let error = match crate::Cli::try_parse_from(args) { + Ok(_) => panic!("file-only flags should require --out"), + Err(error) => error, + }; + assert!(error.to_string().contains("--out")); + } +} + // ------------------------------------------------------------------------- // Notebook discovery // -------------------------------------------------------------------------