metrale_bench/benchmarks/mlperf_agentic/
dataset.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The MLPerf Agentic Inference trajectory file: flat JSONL, one
4//! row per message, each conversation's rows contiguous and numbered 1..=N in
5//! file order.
6//!
7//! Teacher forcing: the prompt for a client turn is the recorded history
8//! (system prompt and recorded user, tool and assistant rows, including
9//! recorded `reasoning_content` and `tool_calls`) plus the current client
10//! message, prebuilt at load time. The model's own output is never fed into
11//! later turns, so per-turn prompts depend only on the file.
12//!
13//! Owner: bench, mlperf_agentic.
14//! Invariants: `load` returns `Ok` only for a file whose every conversation
15//! passes the id, contiguity, turn-number and role checks, and a non-empty
16//! draw.
17
18use std::collections::BTreeSet;
19use std::path::Path;
20
21use anyhow::{Context, Result, bail};
22use serde_json::{Map, Value, json};
23use sha2::{Digest, Sha256};
24
25use super::scoring::{Domain, domain_of};
26
27/// 2026-09-26: One client turn: what is sent, and the recorded assistant row
28/// that scores it.
29#[derive(Clone, Debug)]
30pub struct ClientTurn {
31    /// 2026-09-26: The client row's turn number.
32    pub turn: i64,
33    /// 2026-09-26: Prebuilt prompt: system + recorded history + this client
34    /// message.
35    pub messages: Vec<Value>,
36    /// 2026-09-26: The recorded assistant row that immediately follows this
37    /// client row, raw. `None` for a trailing client row with no assistant
38    /// after it.
39    pub ground_truth: Option<Value>,
40}
41
42#[derive(Clone, Debug)]
43pub struct Conversation {
44    pub id: String,
45    pub domain: Domain,
46    /// 2026-09-26: Tool schemas from the conversation's first `user` row,
47    /// sent with every request of the conversation.
48    pub tools: Option<Value>,
49    pub client_turns: Vec<ClientTurn>,
50}
51
52/// 2026-09-26: How many whole trajectories to take per domain, the first K in
53/// file order, with no RNG. `0` means all of that domain. Conversations are
54/// drawn whole.
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56pub struct DrawSpec {
57    pub coding: usize,
58    pub workflow: usize,
59}
60
61impl DrawSpec {
62    pub fn all() -> Self {
63        Self {
64            coding: 0,
65            workflow: 0,
66        }
67    }
68}
69
70/// 2026-09-26: Load the trajectory file and apply the draw.
71pub fn load(path: &Path, spec: &DrawSpec) -> Result<Vec<Conversation>> {
72    let text =
73        std::fs::read_to_string(path).with_context(|| format!("reading {}", path.display()))?;
74    let conversations = parse(&text)?;
75    let mut taken = Vec::new();
76    let (mut coding_taken, mut workflow_taken) = (0usize, 0usize);
77    for conv in conversations {
78        let (count, limit) = match conv.domain {
79            Domain::Coding => (&mut coding_taken, spec.coding),
80            Domain::Workflow => (&mut workflow_taken, spec.workflow),
81        };
82        if limit == 0 || *count < limit {
83            *count += 1;
84            taken.push(conv);
85        }
86    }
87    if taken.is_empty() {
88        bail!("the draw selected no trajectories");
89    }
90    Ok(taken)
91}
92
93/// 2026-09-26: The draw fingerprint: SHA256 over the ordered conversations'
94/// ids, tools, and each client turn's number, messages and ground truth. Two
95/// runs with equal fingerprints replayed the same prompts, tools and ground
96/// truth in the same order; a differing fingerprint means the numbers are not
97/// comparable, whatever the sample counts say.
98pub fn draw_fingerprint(conversations: &[Conversation]) -> String {
99    let mut digest = Sha256::new();
100    let update = |digest: &mut Sha256, bytes: &[u8]| {
101        digest.update((bytes.len() as u64).to_le_bytes());
102        digest.update(bytes);
103    };
104    for conv in conversations {
105        digest.update(b"conversation\0");
106        update(&mut digest, conv.id.as_bytes());
107        update(
108            &mut digest,
109            &serde_json::to_vec(&conv.tools).expect("JSON values serialize"),
110        );
111        for turn in &conv.client_turns {
112            digest.update(b"turn\0");
113            digest.update(turn.turn.to_le_bytes());
114            update(
115                &mut digest,
116                &serde_json::to_vec(&turn.messages).expect("JSON values serialize"),
117            );
118            update(
119                &mut digest,
120                &serde_json::to_vec(&turn.ground_truth).expect("JSON values serialize"),
121            );
122        }
123    }
124    metrale_closure::hex_lower(&digest.finalize())
125}
126
127fn parse(text: &str) -> Result<Vec<Conversation>> {
128    // 2026-09-26: Group contiguous rows by conversation (non-empty
129    // printable-ASCII ids, rows contiguous); `build_conversation` checks the
130    // turn numbers and the role grammar.
131    let mut order: Vec<String> = Vec::new();
132    let mut groups: std::collections::BTreeMap<String, Vec<Map<String, Value>>> =
133        Default::default();
134    let mut seen_closed: BTreeSet<String> = BTreeSet::new();
135    let mut last_id: Option<String> = None;
136
137    for (i, line) in text.lines().enumerate() {
138        if line.trim().is_empty() {
139            continue;
140        }
141        let row: Map<String, Value> =
142            serde_json::from_str(line).with_context(|| format!("line {}: malformed row", i + 1))?;
143        let id = match row.get("conversation_id") {
144            Some(Value::String(s)) if !s.is_empty() => s.clone(),
145            other => bail!(
146                "line {}: conversation_id must be a non-empty string, got {}",
147                i + 1,
148                other
149                    .map(Value::to_string)
150                    .unwrap_or_else(|| "nothing".into())
151            ),
152        };
153        if id.bytes().any(|b| !(32..=126).contains(&b)) {
154            bail!(
155                "line {}: conversation_id {id:?} has non-printable or non-ASCII characters",
156                i + 1
157            );
158        }
159        if last_id.as_ref() != Some(&id) {
160            if seen_closed.contains(&id) {
161                bail!(
162                    "rows for conversation {id:?} are not consecutive — the loader takes the \
163                     first K per domain in file order, so interleaving would silently change \
164                     which trajectories a draw selects"
165                );
166            }
167            if let Some(prev) = last_id.replace(id.clone()) {
168                seen_closed.insert(prev);
169            }
170            order.push(id.clone());
171        }
172        groups.entry(id).or_default().push(row);
173    }
174
175    let mut out = Vec::new();
176    for id in order {
177        let rows = groups.remove(&id).expect("grouped above");
178        out.push(build_conversation(id, rows)?);
179    }
180    if out.is_empty() {
181        bail!("the trajectory file holds no rows");
182    }
183    Ok(out)
184}
185
186fn turn_of(row: &Map<String, Value>) -> Result<i64> {
187    row.get("turn")
188        .and_then(Value::as_i64)
189        .context("row is missing an integer 'turn'")
190}
191
192fn build_conversation(id: String, rows: Vec<Map<String, Value>>) -> Result<Conversation> {
193    for row in &rows {
194        turn_of(row).with_context(|| format!("conversation {id:?}"))?;
195    }
196    let turns: Vec<i64> = rows.iter().map(|r| turn_of(r).unwrap()).collect();
197    if turns.iter().enumerate().any(|(i, &t)| t != i as i64 + 1) {
198        bail!(
199            "conversation {id:?}: turn numbers must be exactly 1..={} in file order, got \
200             {turns:?}",
201            rows.len()
202        );
203    }
204    validate_roles(&id, &rows)?;
205
206    let system = rows
207        .iter()
208        .find_map(|r| {
209            r.get("system")
210                .and_then(Value::as_str)
211                .filter(|s| !s.is_empty())
212        })
213        .map(str::to_string);
214    let tools = rows
215        .iter()
216        .find(|r| r.get("role").and_then(Value::as_str) == Some("user"))
217        .and_then(|r| r.get("tools"))
218        .filter(|t| !t.is_null())
219        .cloned();
220
221    // 2026-09-26: Single pass in turn order, carrying the running history.
222    // Client rows snapshot (history + current) before extending; assistant
223    // rows only extend.
224    let mut history: Vec<Value> = Vec::new();
225    if let Some(system) = &system {
226        history.push(json!({"role": "system", "content": system}));
227    }
228    let mut client_turns: Vec<ClientTurn> = Vec::new();
229    for (i, row) in rows.iter().enumerate() {
230        let role = row.get("role").and_then(Value::as_str).unwrap_or_default();
231        let row_msgs = format_row(&id, row)?;
232        if role == "user" || role == "tool" {
233            let mut messages = history.clone();
234            messages.extend(row_msgs.iter().cloned());
235            // 2026-09-26: The scorer's ground-truth pairing: the assistant row
236            // that immediately follows this client row in turn order.
237            let ground_truth = rows.get(i + 1).and_then(|next| {
238                (next.get("role").and_then(Value::as_str) == Some("assistant"))
239                    .then(|| Value::Object(next.clone()))
240            });
241            client_turns.push(ClientTurn {
242                turn: turn_of(row).unwrap(),
243                messages,
244                ground_truth,
245            });
246        }
247        history.extend(row_msgs);
248    }
249
250    Ok(Conversation {
251        domain: domain_of(&id),
252        id,
253        tools,
254        client_turns,
255    })
256}
257
258/// 2026-09-26: Format one row into OpenAI message(s): a non-empty
259/// `tool_results` list expands to one tool message per result; otherwise the
260/// recorded message fields pass through, with `content: null` filled in on
261/// tool-calling assistant rows.
262fn format_row(id: &str, row: &Map<String, Value>) -> Result<Vec<Value>> {
263    if let Some(results) = row.get("tool_results").and_then(Value::as_array)
264        && !results.is_empty()
265    {
266        let mut msgs = Vec::new();
267        for (i, result) in results.iter().enumerate() {
268            let (Some(tool_call_id), Some(content)) =
269                (result.get("tool_call_id"), result.get("content"))
270            else {
271                bail!(
272                    "conversation {id:?} turn {}: tool_results[{i}] needs tool_call_id and content",
273                    turn_of(row).unwrap_or_default()
274                );
275            };
276            msgs.push(json!({
277                "role": "tool",
278                "tool_call_id": tool_call_id,
279                "content": content,
280            }));
281        }
282        return Ok(msgs);
283    }
284    let mut msg = Map::new();
285    for key in [
286        "role",
287        "content",
288        "name",
289        "tool_calls",
290        "tool_results",
291        "reasoning_content",
292    ] {
293        if let Some(v) = row.get(key).filter(|v| !v.is_null()) {
294            msg.insert(key.to_string(), v.clone());
295        }
296    }
297    if msg.get("role").and_then(Value::as_str) == Some("assistant")
298        && msg.contains_key("tool_calls")
299        && !msg.contains_key("content")
300    {
301        msg.insert("content".into(), Value::Null);
302    }
303    Ok(if msg.contains_key("role") {
304        vec![Value::Object(msg)]
305    } else {
306        Vec::new()
307    })
308}
309
310/// 2026-09-26: The role grammar: a conversation starts with `user`; `user` is
311/// answered by `assistant`; `assistant` is followed by `tool` or `user`;
312/// `tool` by `assistant` or `user`. A `tool` row that follows an assistant
313/// must follow one that made tool calls, and every `tool` row must carry a
314/// non-empty `tool_results` list.
315fn validate_roles(id: &str, rows: &[Map<String, Value>]) -> Result<()> {
316    let mut state = "start";
317    let mut prev_assistant_had_tool_calls = false;
318    for row in rows {
319        let role = row
320            .get("role")
321            .and_then(Value::as_str)
322            .with_context(|| format!("conversation {id:?}: row without a role"))?;
323        let valid: &[&str] = match state {
324            "start" => &["user"],
325            "user" => &["assistant"],
326            "assistant" => &["tool", "user"],
327            "tool" => &["assistant", "user"],
328            _ => unreachable!(),
329        };
330        if !valid.contains(&role) {
331            bail!(
332                "conversation {id:?} turn {}: got role {role:?} after {state:?}",
333                turn_of(row).unwrap_or_default()
334            );
335        }
336        if role == "tool" {
337            if state == "assistant" && !prev_assistant_had_tool_calls {
338                bail!(
339                    "conversation {id:?} turn {}: a tool row must follow an assistant row \
340                     with non-empty tool_calls",
341                    turn_of(row).unwrap_or_default()
342                );
343            }
344            if row
345                .get("tool_results")
346                .and_then(Value::as_array)
347                .is_none_or(|r| r.is_empty())
348            {
349                bail!(
350                    "conversation {id:?} turn {}: tool rows need a non-empty tool_results list",
351                    turn_of(row).unwrap_or_default()
352                );
353            }
354        }
355        if role == "assistant" {
356            prev_assistant_had_tool_calls = row
357                .get("tool_calls")
358                .and_then(Value::as_array)
359                .is_some_and(|c| !c.is_empty());
360        }
361        state = match role {
362            "user" => "user",
363            "assistant" => "assistant",
364            _ => "tool",
365        };
366    }
367    Ok(())
368}
369
370#[cfg(test)]
371#[path = "dataset_tests.rs"]
372mod tests;
metrale_bench/benchmarks/mlperf_agentic/dataset.rs

metrale_bench/benchmarks/mlperf_agentic/
dataset.rs
