helix/dna/map/
reasoning.rs

1use anyhow::Result;
2use serde::{Deserialize, Serialize};
3use std::path::Path;
4#[derive(Debug, Clone, Serialize, Deserialize)]
5pub struct Message {
6    pub content: String,
7    pub role: String,
8}
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct ReasoningEntry {
11    pub user: String,
12    pub reasoning: String,
13    pub assistant: String,
14    pub template: String,
15    pub conversations: Vec<Message>,
16}
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct ReasoningDataset {
19    pub entries: Vec<ReasoningEntry>,
20}
21impl ReasoningDataset {
22    #[must_use]
23    pub fn new() -> Self {
24        Self { entries: Vec::new() }
25    }
26    pub async fn load<P: AsRef<Path>>(path: P) -> Result<Self> {
27        let content = tokio::fs::read_to_string(path).await?;
28        let dataset: ReasoningDataset = serde_json::from_str(&content)?;
29        Ok(dataset)
30    }
31    pub async fn save<P: AsRef<Path>>(&self, path: P) -> Result<()> {
32        let content = serde_json::to_string_pretty(self)?;
33        tokio::fs::write(path, content).await?;
34        Ok(())
35    }
36    pub fn add_entry(&mut self, entry: ReasoningEntry) {
37        self.entries.push(entry);
38    }
39    #[must_use]
40    pub fn len(&self) -> usize {
41        self.entries.len()
42    }
43    #[must_use]
44    pub fn is_empty(&self) -> bool {
45        self.entries.is_empty()
46    }
47    #[must_use]
48    pub fn create_template(user: &str, reasoning: &str, assistant: &str) -> String {
49        format!(
50            "<|im_start|>user\n{user}<|im_end|>\n<|im_start|>reasoning\n{reasoning}<|im_end|>\n<|im_start|>assistant\n{assistant}<|im_end|>",
51        )
52    }
53}
54impl Default for ReasoningDataset {
55    fn default() -> Self {
56        Self::new()
57    }
58}
59#[cfg(test)]
60mod tests {
61    use super::*;
62    use tempfile::NamedTempFile;
63    #[tokio::test]
64    async fn test_dataset_operations() -> Result<()> {
65        let mut dataset = ReasoningDataset::new();
66        let entry = ReasoningEntry {
67            user: "What motivates Luna?".to_string(),
68            reasoning: "Luna's motivations can be analyzed...".to_string(),
69            assistant: "Luna is motivated by acceptance and self-expression."
70                .to_string(),
71            template: ReasoningDataset::create_template(
72                "What motivates Luna?",
73                "Luna's motivations can be analyzed...",
74                "Luna is motivated by acceptance and self-expression.",
75            ),
76            conversations: vec![
77                Message { content : "What motivates Luna?".to_string(), role : "user"
78                .to_string(), }, Message { content :
79                "Luna's motivations can be analyzed...".to_string(), role : "reasoning"
80                .to_string(), }, Message { content :
81                "Luna is motivated by acceptance and self-expression.".to_string(), role
82                : "assistant".to_string(), },
83            ],
84        };
85        dataset.add_entry(entry);
86        assert_eq!(dataset.len(), 1);
87        let temp_file = NamedTempFile::new()?;
88        dataset.save(temp_file.path()).await?;
89        let loaded_dataset = ReasoningDataset::load(temp_file.path()).await?;
90        assert_eq!(loaded_dataset.len(), 1);
91        Ok(())
92    }
93}