Skip to content

Commit ee326e5

Browse files
committed
feat(llm-obs): add datasets records-all to page past the preview cap
`pup llm-obs datasets records` posts to /api/unstable/llm-obs-mcp/v1/dataset/records, which trims its response to a size budget and returns no cursor. On a dataset whose records carry sizeable inputs that lands at 19 records with `truncated: true`, and because the endpoint emits no cursor there is no way to reach the remainder — `--cursor` exists as a flag but nothing ever produces a value for it. Confirmed against the API directly: --limit 25 and --limit 50 both return 19, and no cursor key appears in the response at any limit. `records-full` is not a workaround either; it caps at 3 ids per call ("record_ids has N entries; cap is 3") and you would still need the full id list to use it. The plain REST route has no such cap and pages properly: GET /api/unstable/llm-obs/v1/datasets/{id}/records?page[limit]=N&page[cursor]=<meta.after> This adds `datasets records-all`, which pages that route via meta.after and emits the aggregate under the same records/returned keys plus pages_fetched. Verified against a 50-record dataset: --limit 7 returns all 50 across 8 pages with 50 unique ids and no __nested_object__ placeholders, where the existing command returns 19. Added as a separate subcommand rather than changing `records` in place, because the two routes return different shapes — `records` gives size-trimmed previews with __nested_object__ / __nested_array__ placeholders, the REST route gives full records. Switching routes under the existing flag would silently change its output shape and break callers that parse it. `records` stays the cheap browse path; this is the explicit "give me everything" path. It takes only --dataset-id, since the REST route does not require a project id. Loop safety: terminates on an empty cursor or a cursor identical to the previous one (a server returning a fixed cursor would otherwise spin), with a 200-page hard stop that warns on stdout before returning what it has. Tests cover the happy path (asserts BOTH pages were requested, proving the cursor is followed), the repeated-cursor guard, and a 500 response.
1 parent df43130 commit ee326e5

3 files changed

Lines changed: 176 additions & 2 deletions

File tree

docs/COMMANDS.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ pup <domain> <subgroup> <action> [options] # Nested commands
6666
| data-deletion | requests (list, create, cancel) | src/commands/data_deletion.rs ||
6767
| data-governance | scanner-rules (list) | src/commands/data_governance.rs ||
6868
| obs-pipelines | list, get, create, update, delete, validate | src/commands/obs_pipelines.rs ||
69-
| 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 ||
69+
| 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-all, records-full), spans (search) | src/commands/llm_obs.rs ||
7070
| reference-tables | list, get, create, batch-query | src/commands/reference_tables.rs ||
7171
| network | flows list, devices (list, get, interfaces, tags), interfaces (list, update) | src/commands/network.rs ||
7272
| cloud | aws, gcp, azure, oci | src/commands/cloud.rs ||
@@ -300,7 +300,7 @@ steps) bypass `format_and_print` and do not honor `--jq`.
300300

301301
### v0.28.0 — New Command Groups and Full Pipeline Implementation
302302

303-
-**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)
303+
-**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-all, records-full), spans (search)
304304
-**reference-tables** (new) — Reference table management (list, get, create, batch-query)
305305
-**obs-pipelines** (upgraded from placeholder) — Full CRUD: list, get, create, update, delete, validate
306306
- **costs** — Added cloud cost configs: `aws-config`, `azure-config`, `gcp-config` (list, get, create, delete each)

src/commands/llm_obs.rs

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,71 @@ pub async fn datasets_records(
200200
formatter::output(cfg, &resp)
201201
}
202202

203+
/// Hard stop on page count so a misbehaving cursor cannot loop forever.
204+
const RECORDS_ALL_MAX_PAGES: u32 = 200;
205+
206+
/// Read EVERY record in a dataset by paging the REST records endpoint.
207+
///
208+
/// `datasets records` posts to `llm-obs-mcp/v1/dataset/records`, which trims its response to a
209+
/// size budget (~19 records on a dataset with sizeable inputs), reports `truncated: true`, and
210+
/// returns **no cursor** — so there is no way to reach the rest of a large dataset through it.
211+
/// This pages `GET /api/unstable/llm-obs/v1/datasets/{id}/records` using `meta.after`, which has
212+
/// no such cap, and emits the aggregated records under the same `records`/`returned` keys.
213+
///
214+
/// The REST route needs no project_id, so this takes only the dataset id.
215+
pub async fn datasets_records_all(
216+
cfg: &Config,
217+
dataset_id: &str,
218+
limit: Option<u32>,
219+
) -> Result<()> {
220+
let path = format!("/api/unstable/llm-obs/v1/datasets/{dataset_id}/records");
221+
let page_limit = limit.unwrap_or(100).to_string();
222+
let mut records: Vec<serde_json::Value> = Vec::new();
223+
let mut cursor = String::new();
224+
let mut pages: u32 = 0;
225+
226+
loop {
227+
let mut query: Vec<(&str, &str)> = vec![("page[limit]", page_limit.as_str())];
228+
if !cursor.is_empty() {
229+
query.push(("page[cursor]", cursor.as_str()));
230+
}
231+
let resp = raw_client::raw_get(cfg, &path, &query)
232+
.await
233+
.map_err(|e| anyhow::anyhow!("failed to list dataset records: {e:?}"))?;
234+
pages += 1;
235+
236+
match resp["data"].as_array() {
237+
Some(page) if !page.is_empty() => records.extend(page.iter().cloned()),
238+
_ => break,
239+
}
240+
241+
let after = resp["meta"]["after"].as_str().unwrap_or_default();
242+
// Empty or unchanged cursor both mean "no further pages".
243+
if after.is_empty() || after == cursor {
244+
break;
245+
}
246+
cursor = after.to_string();
247+
248+
if pages >= RECORDS_ALL_MAX_PAGES {
249+
eprintln!(
250+
"warning: stopped after {pages} pages ({} records); the dataset may have more",
251+
records.len()
252+
);
253+
break;
254+
}
255+
}
256+
257+
let out = serde_json::json!({
258+
"dataset_id": dataset_id,
259+
"kind": "full_list",
260+
"records": records,
261+
"returned": records.len(),
262+
"truncated": false,
263+
"pages_fetched": pages,
264+
});
265+
formatter::output(cfg, &out)
266+
}
267+
203268
pub async fn datasets_records_full(
204269
cfg: &Config,
205270
project_id: &str,
@@ -1029,6 +1094,105 @@ mod tests {
10291094
std::env::remove_var("DD_TOKEN_STORAGE");
10301095
}
10311096

1097+
#[tokio::test]
1098+
async fn test_llm_obs_datasets_records_all_pages_until_cursor_empty() {
1099+
let _lock = lock_env().await;
1100+
std::env::set_var("DD_TOKEN_STORAGE", "file");
1101+
let mut server = mockito::Server::new_async().await;
1102+
let cfg = test_config(&server.url());
1103+
1104+
// Page 1: two records plus a cursor. Matched by the ABSENCE of page[cursor].
1105+
let page1 = server
1106+
.mock("GET", mockito::Matcher::Any)
1107+
.match_query(mockito::Matcher::UrlEncoded(
1108+
"page[limit]".into(),
1109+
"2".into(),
1110+
))
1111+
.with_status(200)
1112+
.with_header("content-type", "application/json")
1113+
.with_body(r#"{"data":[{"id":"rec-1"},{"id":"rec-2"}],"meta":{"after":"CURSOR2"}}"#)
1114+
.expect(1)
1115+
.create_async()
1116+
.await;
1117+
1118+
// Page 2: one record and an empty cursor, which terminates the loop.
1119+
let page2 = server
1120+
.mock("GET", mockito::Matcher::Any)
1121+
.match_query(mockito::Matcher::UrlEncoded(
1122+
"page[cursor]".into(),
1123+
"CURSOR2".into(),
1124+
))
1125+
.with_status(200)
1126+
.with_header("content-type", "application/json")
1127+
.with_body(r#"{"data":[{"id":"rec-3"}],"meta":{"after":""}}"#)
1128+
.expect(1)
1129+
.create_async()
1130+
.await;
1131+
1132+
let result = super::datasets_records_all(&cfg, "ds-1", Some(2)).await;
1133+
assert!(
1134+
result.is_ok(),
1135+
"datasets_records_all failed: {:?}",
1136+
result.err()
1137+
);
1138+
// Both pages must have been requested — proves the cursor was followed.
1139+
page1.assert_async().await;
1140+
page2.assert_async().await;
1141+
1142+
cleanup_env();
1143+
std::env::remove_var("DD_TOKEN_STORAGE");
1144+
}
1145+
1146+
#[tokio::test]
1147+
async fn test_llm_obs_datasets_records_all_stops_on_repeated_cursor() {
1148+
let _lock = lock_env().await;
1149+
std::env::set_var("DD_TOKEN_STORAGE", "file");
1150+
let mut server = mockito::Server::new_async().await;
1151+
let cfg = test_config(&server.url());
1152+
1153+
// Always returns the same cursor it was given: a server bug that would loop forever
1154+
// if the guard were missing.
1155+
let _mock = mock_any(
1156+
&mut server,
1157+
"GET",
1158+
r#"{"data":[{"id":"rec-1"}],"meta":{"after":"SAME"}}"#,
1159+
)
1160+
.await;
1161+
1162+
let result = super::datasets_records_all(&cfg, "ds-1", None).await;
1163+
assert!(result.is_ok(), "expected ok, got {:?}", result.err());
1164+
1165+
cleanup_env();
1166+
std::env::remove_var("DD_TOKEN_STORAGE");
1167+
}
1168+
1169+
#[tokio::test]
1170+
async fn test_llm_obs_datasets_records_all_500() {
1171+
let _lock = lock_env().await;
1172+
std::env::set_var("DD_TOKEN_STORAGE", "file");
1173+
let mut server = mockito::Server::new_async().await;
1174+
let cfg = test_config(&server.url());
1175+
1176+
let _mock = server
1177+
.mock("GET", mockito::Matcher::Any)
1178+
.match_query(mockito::Matcher::Any)
1179+
.with_status(500)
1180+
.with_header("content-type", "application/json")
1181+
.with_body(r#"{"errors":["internal error"]}"#)
1182+
.create_async()
1183+
.await;
1184+
1185+
let result = super::datasets_records_all(&cfg, "ds-1", None).await;
1186+
assert!(
1187+
result.is_err(),
1188+
"expected error but got ok: {:?}",
1189+
result.ok()
1190+
);
1191+
1192+
cleanup_env();
1193+
std::env::remove_var("DD_TOKEN_STORAGE");
1194+
}
1195+
10321196
#[tokio::test]
10331197
async fn test_llm_obs_experiments_create() {
10341198
let _lock = lock_env().await;

src/main.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9410,6 +9410,13 @@ enum LlmObsDatasetsActions {
94109410
#[arg(long, help = "JSON file with restore version body (required)")]
94119411
file: String,
94129412
},
9413+
/// Read ALL dataset records, paging past the preview endpoint's response-size cap
9414+
RecordsAll {
9415+
#[arg(long, help = "Dataset ID (required)")]
9416+
dataset_id: String,
9417+
#[arg(long, help = "Records per page (default 100)")]
9418+
limit: Option<u32>,
9419+
},
94139420
/// Read dataset records (structure-preserving previews + schema summary)
94149421
Records {
94159422
#[arg(long, help = "Project ID (required)")]
@@ -16261,6 +16268,9 @@ async fn main_inner() -> anyhow::Result<()> {
1626116268
commands::llm_obs::datasets_restore(&cfg, &project_id, &dataset_id, &file)
1626216269
.await?;
1626316270
}
16271+
LlmObsDatasetsActions::RecordsAll { dataset_id, limit } => {
16272+
commands::llm_obs::datasets_records_all(&cfg, &dataset_id, limit).await?;
16273+
}
1626416274
LlmObsDatasetsActions::Records {
1626516275
project_id,
1626616276
dataset_id,

0 commit comments

Comments
 (0)