use crate::build::{Msg, Row};
use crate::clean::clean_text;
use crate::config::SourceConfig;
use serde_json::Value;
use std::path::{Path, PathBuf};
pub struct DatasetLoad {
pub rows: Vec<(String, Row)>,
pub resolved_path: PathBuf,
pub skipped: usize,
}
fn resolve_path(raw: &Path) -> std::io::Result<PathBuf> {
if raw.is_file() {
return Ok(raw.to_path_buf());
}
if raw.is_dir() {
let mut best: Option<(u64, PathBuf)> = None;
for entry in std::fs::read_dir(raw)? {
let path = entry?.path();
if path.extension().and_then(|e| e.to_str()) == Some("jsonl") {
let size = path.metadata().map(|m| m.len()).unwrap_or(0);
if best.as_ref().map(|(s, _)| size > *s).unwrap_or(true) {
best = Some((size, path));
}
}
}
if let Some((_, p)) = best {
return Ok(p);
}
}
Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("no dataset .jsonl found at {}", raw.display()),
))
}
fn normalize_row(
value: &Value,
system_prompt: Option<&str>,
user_prefix: Option<&str>,
preserve: bool,
) -> Option<Row> {
let messages = value.get("messages")?.as_array()?;
if messages.len() < 2 {
return None;
}
let mut kept: Vec<Msg> = Vec::new();
for m in messages {
let role = m.get("role").and_then(|r| r.as_str()).unwrap_or("");
if !matches!(role, "system" | "user" | "assistant") {
continue;
}
let content = match m.get("content").and_then(|c| c.as_str()) {
Some(s) => s,
None => continue,
};
let cleaned = if preserve { content.to_string() } else { clean_text(content) };
kept.push(Msg { role: role.to_string(), content: cleaned });
}
if kept.is_empty() || kept.last().map(|m| m.role.as_str()) != Some("assistant") {
return None;
}
if let Some(prefix) = user_prefix {
if let Some(first_user) = kept.iter_mut().find(|m| m.role == "user") {
first_user.content = format!("{prefix}{}", first_user.content);
}
}
if let Some(sys) = system_prompt {
if kept.first().map(|m| m.role.as_str()) == Some("system") {
kept[0].content = sys.to_string();
} else {
kept.insert(0, Msg { role: "system".to_string(), content: sys.to_string() });
}
}
Some(Row { messages: kept })
}
pub fn load_dataset_rows(source: &SourceConfig) -> std::io::Result<DatasetLoad> {
let resolved_path = resolve_path(Path::new(&source.path))?;
let text = std::fs::read_to_string(&resolved_path)?;
let name = source.source_name();
let mut rows = Vec::new();
let mut skipped = 0usize;
let mut index = 0usize;
for line in text.lines() {
let line = line.trim();
if line.is_empty() {
continue;
}
if let Some(cap) = source.max_rows {
if rows.len() >= cap {
break;
}
}
let value: Value = match serde_json::from_str(line) {
Ok(v) => v,
Err(_) => {
skipped += 1;
continue;
}
};
let row_key = value
.get("id")
.and_then(|v| v.as_str().map(|s| s.to_string()).or_else(|| v.as_i64().map(|n| n.to_string())))
.unwrap_or_else(|| index.to_string());
let preserve = source.preserve_code.unwrap_or(false);
match normalize_row(&value, source.system_prompt.as_deref(), source.user_prefix.as_deref(), preserve) {
Some(row) => rows.push((format!("{name}:{row_key}"), row)),
None => skipped += 1,
}
index += 1;
}
Ok(DatasetLoad { rows, resolved_path, skipped })
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
fn src(path: &str, name: &str, system: Option<&str>, prefix: Option<&str>) -> SourceConfig {
SourceConfig {
path: path.to_string(),
r#type: Some("dataset".to_string()),
name: Some(name.to_string()),
max_rows: None,
system_prompt: system.map(|s| s.to_string()),
user_prefix: prefix.map(|s| s.to_string()),
role: None,
preserve_code: None, refresh: None,
}
}
#[test]
fn normalizes_injects_and_wraps() {
let dir = std::env::temp_dir().join(format!("kibble_ds_{}", std::process::id()));
let _ = fs::remove_dir_all(&dir);
fs::create_dir_all(&dir).unwrap();
let p = dir.join("d.jsonl");
let lines = [
r#"{"id":"a","messages":[{"role":"user","content":"hi @bob"},{"role":"assistant","content":"yo"}]}"#,
r#"{"messages":[{"role":"user","content":"only user"}]}"#,
r#"{not json}"#,
];
fs::write(&p, lines.join("\n")).unwrap();
let load = load_dataset_rows(&src(
p.to_str().unwrap(),
"ds",
Some("SYS"),
Some("PFX "),
))
.unwrap();
assert_eq!(load.rows.len(), 1);
assert_eq!(load.skipped, 2);
let (id, row) = &load.rows[0];
assert_eq!(id, "ds:a");
assert_eq!(row.messages[0].role, "system");
assert_eq!(row.messages[0].content, "SYS");
assert_eq!(row.messages[1].role, "user");
assert_eq!(row.messages[1].content, "PFX hi");
assert_eq!(row.messages[2].content, "yo");
}
#[test]
fn resolves_largest_jsonl_in_dir() {
let dir = std::env::temp_dir().join(format!("kibble_dsdir_{}", std::process::id()));
let _ = fs::remove_dir_all(&dir);
fs::create_dir_all(&dir).unwrap();
fs::write(dir.join("small.jsonl"), "x").unwrap();
fs::write(
dir.join("big.jsonl"),
r#"{"messages":[{"role":"user","content":"q"},{"role":"assistant","content":"a"}]}"#,
)
.unwrap();
let load = load_dataset_rows(&src(dir.to_str().unwrap(), "d", None, None)).unwrap();
assert!(load.resolved_path.ends_with("big.jsonl"));
assert_eq!(load.rows.len(), 1);
}
#[test]
fn missing_path_errors() {
assert!(load_dataset_rows(&src("/no/such.jsonl", "d", None, None)).is_err());
}
#[test]
fn skips_non_string_content_turn_not_whole_row() {
let dir = std::env::temp_dir()
.join(format!("kibble_ds_nonstr_{}", std::process::id()));
let _ = fs::remove_dir_all(&dir);
fs::create_dir_all(&dir).unwrap();
let p = dir.join("d.jsonl");
let line = r#"{"id":"m","messages":[{"role":"user","content":"hello"},{"role":"system","content":["multi","part"]},{"role":"assistant","content":"hi back"}]}"#;
fs::write(&p, line).unwrap();
let load = load_dataset_rows(&src(p.to_str().unwrap(), "ds", None, None)).unwrap();
assert_eq!(load.rows.len(), 1, "row should be kept");
assert_eq!(load.skipped, 0);
let (_, row) = &load.rows[0];
assert_eq!(row.messages.len(), 2, "only user + assistant kept");
assert_eq!(row.messages[0].role, "user");
assert_eq!(row.messages[1].role, "assistant");
}
}