navi_core/
dataset.rs

1//! Trace-to-dataset export for evals, routing, and permission tuning.
2
3use crate::eval::{EvalCase, eval_case_from_trace};
4use crate::security::redact_secrets;
5use crate::trace::{TurnOutcome, TurnTrace};
6use anyhow::{Context, Result};
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10pub struct DatasetRow {
11    pub version: u32,
12    pub row_type: DatasetRowType,
13    pub task: String,
14    pub outcome: String,
15    pub tools: Vec<String>,
16    #[serde(default)]
17    pub capabilities: Vec<String>,
18    #[serde(default)]
19    pub mcp_tainted: bool,
20    pub verifier_passed: bool,
21    pub reward: f64,
22}
23
24#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
25#[serde(rename_all = "snake_case")]
26pub enum DatasetRowType {
27    PreferencePair,
28    NegativeExample,
29    ToolRouterTraining,
30    PermissionClassifier,
31    VerifierReward,
32}
33
34pub fn trace_to_dataset_rows(trace: &TurnTrace) -> Vec<DatasetRow> {
35    let tools = trace
36        .tool_calls
37        .iter()
38        .map(|call| call.invocation.tool_name.clone())
39        .collect::<Vec<_>>();
40    let verifier_passed = trace
41        .verifier_results
42        .iter()
43        .any(|verifier| verifier.passed);
44    let success = matches!(
45        trace.outcome,
46        TurnOutcome::Success | TurnOutcome::PartialSuccess
47    );
48    let task = redact_secrets(&trace.task);
49    let capabilities = trace
50        .capabilities
51        .iter()
52        .map(|entry| entry.capability.as_key())
53        .collect::<std::collections::BTreeSet<_>>()
54        .into_iter()
55        .collect::<Vec<_>>();
56    let mcp_tainted = trace.tool_calls.iter().any(|call| {
57        call.result
58            .output
59            .get("tainted")
60            .and_then(serde_json::Value::as_bool)
61            .unwrap_or(false)
62            || call
63                .result
64                .output
65                .get("provenance")
66                .and_then(|value| value.get("source"))
67                .and_then(serde_json::Value::as_str)
68                .is_some_and(|source| source == "mcp")
69    });
70    let mut rows = vec![
71        DatasetRow {
72            version: 1,
73            row_type: if success {
74                DatasetRowType::ToolRouterTraining
75            } else {
76                DatasetRowType::NegativeExample
77            },
78            task: task.clone(),
79            outcome: outcome_label(&trace.outcome),
80            tools: tools.clone(),
81            capabilities: capabilities.clone(),
82            mcp_tainted,
83            verifier_passed,
84            reward: if success && verifier_passed { 1.0 } else { 0.0 },
85        },
86        DatasetRow {
87            version: 1,
88            row_type: DatasetRowType::VerifierReward,
89            task,
90            outcome: outcome_label(&trace.outcome),
91            tools,
92            capabilities: capabilities.clone(),
93            mcp_tainted,
94            verifier_passed,
95            reward: if verifier_passed { 1.0 } else { -1.0 },
96        },
97    ];
98    if !capabilities.is_empty() {
99        rows.push(DatasetRow {
100            version: 1,
101            row_type: DatasetRowType::PermissionClassifier,
102            task: redact_secrets(&trace.task),
103            outcome: outcome_label(&trace.outcome),
104            tools: trace
105                .tool_calls
106                .iter()
107                .map(|call| call.invocation.tool_name.clone())
108                .collect(),
109            capabilities,
110            mcp_tainted,
111            verifier_passed,
112            reward: if success { 1.0 } else { -1.0 },
113        });
114    }
115    rows
116}
117
118pub fn traces_to_eval_candidates(traces: &[TurnTrace]) -> Vec<EvalCase> {
119    traces.iter().filter_map(eval_case_from_trace).collect()
120}
121
122pub fn export_jsonl(rows: &[DatasetRow]) -> Result<String> {
123    let mut lines = Vec::with_capacity(rows.len());
124    for row in rows {
125        lines.push(serde_json::to_string(row).context("failed to serialize dataset row to JSON")?);
126    }
127    Ok(lines.join("\n"))
128}
129
130fn outcome_label(outcome: &TurnOutcome) -> String {
131    match outcome {
132        TurnOutcome::Success => "success".to_string(),
133        TurnOutcome::PartialSuccess => "partial_success".to_string(),
134        TurnOutcome::Stopped(reason) => format!("stopped:{reason}"),
135        TurnOutcome::Failed(reason) => format!("failed:{reason}"),
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn dataset_export_redacts_task_secrets() {
145        let mut trace = TurnTrace::new("t", "s", "p", "m", "token=sk-12345678901234567890");
146        trace.outcome = TurnOutcome::Failed("no".to_string());
147
148        let rows = trace_to_dataset_rows(&trace);
149        let jsonl = export_jsonl(&rows).expect("export");
150
151        assert!(!jsonl.contains("sk-1234567890"));
152        assert!(jsonl.contains("negative_example"));
153    }
154
155    #[test]
156    fn dataset_export_includes_permission_classifier_rows() {
157        let mut trace = TurnTrace::new("t", "s", "p", "m", "read repo");
158        trace.record_capability(crate::capability::CapabilityLedgerEntry {
159            capability: crate::capability::Capability::RepoRead,
160            scope: crate::capability::CapabilityScope::Turn("t".to_string()),
161            decision: crate::capability::CapabilityDecision::Consumed,
162            at_ms: 1,
163            justification: "read".to_string(),
164        });
165
166        let rows = trace_to_dataset_rows(&trace);
167
168        assert!(
169            rows.iter()
170                .any(|row| row.row_type == DatasetRowType::PermissionClassifier
171                    && row.capabilities == vec!["repo.read"])
172        );
173    }
174}
navi_core/dataset.rs

navi_core/
dataset.rs
