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

navi_core/
dataset.rs
