sail_scripts/agent_core/src/lib.rs

use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use serde::{Deserialize, Serialize};
use std::fs::File;
use std::io::Read;
use rayon::prelude::*;

#[derive(Serialize, Deserialize, Debug)]
struct Tool {
    name: String,
    args: String,
    result: String,
}

#[derive(Serialize, Deserialize, Debug)]
struct InstructionData {
    instruction: String,
    response: String,
    system: Option<String>,
    thought: Option<String>,
    tools: Option<Vec<Tool>>,
}

/// A highly memory-safe Rust function to parse large JSON arrays into structured format strings.
/// This completely bypasses Python's GIL and memory-overhead bottlenecks.
#[pyfunction]
fn parse_instruction_dataset_safe(file_path: String) -> PyResult<Vec<String>> {
    let mut file = match File::open(&file_path) {
        Ok(f) => f,
        Err(_) => return Err(pyo3::exceptions::PyIOError::new_err(format!("Failed to open {}", file_path))),
    };

    let mut contents = String::new();
    if file.read_to_string(&mut contents).is_err() {
        return Err(pyo3::exceptions::PyIOError::new_err("Failed to read file contents"));
    }

    let parsed_data: Vec<InstructionData> = match serde_json::from_str(&contents) {
        Ok(d) => d,
        Err(_) => return Err(pyo3::exceptions::PyValueError::new_err("Failed to parse JSON. Ensure it is a valid array of instruction objects.")),
    };

    // Use Rayon to parallel-process strings across CPU cores natively 
    let formatted_texts: Vec<String> = parsed_data.into_par_iter().map(|item| {
        let system_prompt = item.system.unwrap_or_else(|| "You are a helpful, smart AI assistant.".to_string());
        
        let thought_str = match item.thought {
            Some(t) => format!("[THOUGHT] {} ", t),
            None => "".to_string(),
        };

        let tools_str = match item.tools {
            Some(tools) => {
                let mut ts = String::new();
                for tool in tools {
                    ts.push_str(&format!("[TOOL_CALL] {} [TOOL_ARG] {} [TOOL_RESULT] {} ", tool.name, tool.args, tool.result));
                }
                ts
            },
            None => "".to_string(),
        };

        CowFormat(system_prompt, item.instruction, thought_str, tools_str, item.response)
    }).collect();

    Ok(formatted_texts)
}

/// Ingests raw text files recursively from a directory (Common Crawl style).
/// Highly performance-optimized for large-scale pre-training data.
#[pyfunction]
fn parse_text_directory(dir_path: String) -> PyResult<Vec<String>> {
    use walkdir::WalkDir;
    
    let mut files = Vec::new();
    for entry in WalkDir::new(dir_path).into_iter().filter_map(|e| e.ok()) {
        if entry.file_type().is_file() {
            if let Some(ext) = entry.path().extension() {
                if ext == "txt" {
                    files.push(entry.path().to_owned());
                }
            }
        }
    }

    let texts: Vec<String> = files.into_par_iter().map(|path| {
        let mut f = File::open(path).expect("Failed to open file");
        let mut content = String::new();
        f.read_to_string(&mut content).unwrap_or(0);
        content
    }).collect();

    Ok(texts)
}

fn CowFormat(sys: String, inst: String, thought: String, tools: String, resp: String) -> String {
    format!("[SYSTEM] {} [USER] {} {}{}[RESPONSE] {} <EOS>", sys, inst, thought, tools, resp)
}

/// The Python Module entry point
#[pymodule]
fn agent_core(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(parse_instruction_dataset_safe, m)?)?;
    m.add_function(wrap_pyfunction!(parse_text_directory, m)?)?;
    Ok(())
}