Merge pull request #674 from gsvigruha/llmobs-add-dataset-records-and-experiment-events

feat(llm-obs): add dataset-records reads + experiment eval-metric submit

Author: platinummonkey

README.md

@@ -180,7 +180,7 @@ list of commands as built.
 | Data Streams (Kafka) | ✅ | `kafka topic-configs`, `kafka broker-configs`, `kafka client-configs`, `kafka read-messages` | **Experimental** — Kafka cluster inspection via Datadog |
 | Restricted Datasets | ✅ | `datasets list`, `datasets get`, `datasets create`, `datasets update`, `datasets delete` | Restricted dataset management for data access control |
 | Observability Pipelines | ✅ | `obs-pipelines list`, `obs-pipelines get`, `obs-pipelines create`, `obs-pipelines update`, `obs-pipelines delete`, `obs-pipelines validate` | Full pipeline CRUD and validation |
-| LLM Observability | ✅ | `llm-obs projects`, `llm-obs experiments`, `llm-obs datasets` | **New** — LLM Obs projects, experiments, and dataset management |
+| LLM Observability | ✅ | `llm-obs projects`, `llm-obs experiments`, `llm-obs datasets` | **New** — LLM Obs projects, experiments (incl. `events submit`), and dataset management (incl. `datasets records` / `records-full`) |
 | Reference Tables | ✅ | `reference-tables list`, `reference-tables get`, `reference-tables create`, `reference-tables batch-query` | **New** — Reference table management for log enrichment |
 | Miscellaneous | ✅ | `misc ip-ranges`, `misc status` | IP ranges and status |
 | App Builder | ✅ | `app-builder list`, `app-builder get`, `app-builder create`, `app-builder update`, `app-builder delete`, `app-builder publish` | Low-code app management with publish/unpublish and batch delete |

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), metric-values, dimension-values), datasets (create, list, batch-update, clone, restore), spans (search) | 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-full), 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 | ✅ |
@@ -300,7 +300,7 @@ steps) bypass `format_and_print` and do not honor `--jq`.
 
 ### 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, batch-update, clone, restore), spans (search)
+- ✅ **llm-obs** (new) — LLM Observability: 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-full), 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)

src/commands/llm_obs.rs

@@ -156,6 +156,71 @@ pub async fn datasets_restore(
     Ok(())
 }
 
+// ---- Dataset records (no typed equivalent — unstable MCP endpoints) ----
+
+#[allow(clippy::too_many_arguments)]
+pub async fn datasets_records(
+    cfg: &Config,
+    project_id: &str,
+    dataset_id: &str,
+    record_ids: Option<Vec<String>>,
+    tags: Option<Vec<String>>,
+    canonical_id: Option<String>,
+    dataset_version: Option<i64>,
+    limit: u32,
+    cursor: Option<String>,
+    compute_schema: Option<bool>,
+) -> Result<()> {
+    let mut body = serde_json::json!({
+        "project_id": project_id,
+        "dataset_id": dataset_id,
+        "limit": limit,
+    });
+    if let Some(ids) = record_ids {
+        body["record_ids"] = serde_json::json!(ids);
+    }
+    if let Some(t) = tags {
+        body["tags"] = serde_json::json!(t);
+    }
+    if let Some(c) = canonical_id {
+        body["canonical_id"] = serde_json::json!(c);
+    }
+    if let Some(v) = dataset_version {
+        body["dataset_version"] = serde_json::json!(v);
+    }
+    if let Some(c) = cursor {
+        body["cursor"] = serde_json::json!(c);
+    }
+    if let Some(cs) = compute_schema {
+        body["compute_schema"] = serde_json::json!(cs);
+    }
+    let resp = raw_client::raw_post(cfg, "/api/unstable/llm-obs-mcp/v1/dataset/records", body)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to get dataset records: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+pub async fn datasets_records_full(
+    cfg: &Config,
+    project_id: &str,
+    dataset_id: &str,
+    record_ids: Vec<String>,
+) -> Result<()> {
+    let body = serde_json::json!({
+        "project_id": project_id,
+        "dataset_id": dataset_id,
+        "record_ids": record_ids,
+    });
+    let resp = raw_client::raw_post(
+        cfg,
+        "/api/unstable/llm-obs-mcp/v1/dataset/records-full",
+        body,
+    )
+    .await
+    .map_err(|e| anyhow::anyhow!("failed to get full dataset records: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
 // ---- Experiment analytics (no typed equivalent — unstable MCP endpoints) ----
 
 pub async fn experiments_summary(cfg: &Config, experiment_id: &str) -> Result<()> {
@@ -214,6 +279,34 @@ pub async fn experiments_events_get(
     formatter::output(cfg, &resp)
 }
 
+pub async fn experiments_events_submit(
+    cfg: &Config,
+    experiment_id: &str,
+    metrics: &str,
+    tags: Option<Vec<String>>,
+) -> Result<()> {
+    // Mirrors the submit_llmobs_experiment_events MCP tool: experiment_id + metrics
+    // (array) + optional tags. metrics is passed through as-is; the server validates it.
+    let metrics: serde_json::Value = serde_json::from_str(metrics).map_err(|e| {
+        anyhow::anyhow!("--metrics must be a JSON array of eval-metric events: {e}")
+    })?;
+    let mut body = serde_json::json!({
+        "experiment_id": experiment_id,
+        "metrics": metrics,
+    });
+    if let Some(t) = tags {
+        body["tags"] = serde_json::json!(t);
+    }
+    let resp = raw_client::raw_post(
+        cfg,
+        "/api/unstable/llm-obs-mcp/v1/experiment/ingest-events",
+        body,
+    )
+    .await
+    .map_err(|e| anyhow::anyhow!("failed to submit experiment events: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
 pub async fn experiments_metric_values(
     cfg: &Config,
     experiment_id: &str,
@@ -2851,4 +2944,189 @@ mod tests {
         cleanup_env();
         std::env::remove_var("DD_TOKEN_STORAGE");
     }
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_records() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let body = r#"{"status":"success","data":{"records":[{"id":"rec-1"}],"schema_summary":{},"returned":1,"truncated":false,"next_cursor":null}}"#;
+        let _mock = mock_post(
+            &mut server,
+            "/api/unstable/llm-obs-mcp/v1/dataset/records",
+            200,
+            body,
+        )
+        .await;
+
+        let result = super::datasets_records(
+            &cfg, "proj-1", "ds-1", None, None, None, None, 10, None, None,
+        )
+        .await;
+        assert!(
+            result.is_ok(),
+            "datasets_records failed: {:?}",
+            result.err()
+        );
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_records_filtered() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let body = r#"{"status":"success","data":{"records":[],"returned":0,"truncated":false,"next_cursor":null}}"#;
+        let _mock = mock_post(
+            &mut server,
+            "/api/unstable/llm-obs-mcp/v1/dataset/records",
+            200,
+            body,
+        )
+        .await;
+
+        let result = super::datasets_records(
+            &cfg,
+            "proj-1",
+            "ds-1",
+            Some(vec!["rec-1".into(), "rec-2".into()]),
+            Some(vec!["env:prod".into()]),
+            Some("canon-1".into()),
+            Some(3),
+            5,
+            Some("cursor-abc".into()),
+            Some(false),
+        )
+        .await;
+        assert!(
+            result.is_ok(),
+            "datasets_records filtered failed: {:?}",
+            result.err()
+        );
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_records_500() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let _mock = mock_post(
+            &mut server,
+            "/api/unstable/llm-obs-mcp/v1/dataset/records",
+            500,
+            r#"{"errors":["internal server error"]}"#,
+        )
+        .await;
+
+        let result = super::datasets_records(
+            &cfg, "proj-1", "ds-1", None, None, None, None, 10, None, None,
+        )
+        .await;
+        assert!(result.is_err(), "should fail on 500");
+        assert!(result.unwrap_err().to_string().contains("500"));
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_records_full() {
+        let _lock = lock_env().await;
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let body = r#"{"status":"success","data":{"records":[{"id":"rec-1","input":{"prompt":"hello"},"expected_output":"world"}]}}"#;
+        let _mock = mock_post(
+            &mut server,
+            "/api/unstable/llm-obs-mcp/v1/dataset/records-full",
+            200,
+            body,
+        )
+        .await;
+
+        let result =
+            super::datasets_records_full(&cfg, "proj-1", "ds-1", vec!["rec-1".into()]).await;
+        assert!(
+            result.is_ok(),
+            "datasets_records_full failed: {:?}",
+            result.err()
+        );
+        cleanup_env();
+    }
+
+    #[tokio::test]
+    async fn test_llm_obs_experiments_events_submit() {
+        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 resp_body = r#"{"status":"success","data":{"accepted":1}}"#;
+        let _mock = mock_post(
+            &mut server,
+            "/api/unstable/llm-obs-mcp/v1/experiment/ingest-events",
+            200,
+            resp_body,
+        )
+        .await;
+
+        let result = super::experiments_events_submit(
+            &cfg,
+            "exp-1",
+            r#"[{"label":"accuracy","metric_type":"score","score_value":0.9}]"#,
+            Some(vec!["run:1".to_string()]),
+        )
+        .await;
+        assert!(
+            result.is_ok(),
+            "experiments_events_submit failed: {:?}",
+            result.err()
+        );
+        cleanup_env();
+        std::env::remove_var("DD_TOKEN_STORAGE");
+    }
+
+    #[tokio::test]
+    async fn test_llm_obs_experiments_events_submit_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 _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::experiments_events_submit(
+            &cfg,
+            "exp-1",
+            r#"[{"label":"accuracy","metric_type":"score","score_value":0.9}]"#,
+            None,
+        )
+        .await;
+        assert!(result.is_err(), "should fail on 400");
+        cleanup_env();
+        std::env::remove_var("DD_TOKEN_STORAGE");
+    }
+
+    #[tokio::test]
+    async fn test_llm_obs_experiments_events_submit_invalid_json() {
+        let _lock = lock_env().await;
+        std::env::set_var("DD_TOKEN_STORAGE", "file");
+        let server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        // Malformed --metrics should fail locally before any request is made.
+        let result = super::experiments_events_submit(&cfg, "exp-1", "not-json", None).await;
+        assert!(result.is_err(), "should fail on invalid metrics JSON");
+        cleanup_env();
+        std::env::remove_var("DD_TOKEN_STORAGE");
+    }
 }