helix/dna/map/
reasoning.rs1use 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}