use super::metrics::SessionMetrics;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EntrenarSession {
pub id: String,
pub name: String,
pub user: Option<String>,
pub created_at: chrono::DateTime<chrono::Utc>,
pub ended_at: Option<chrono::DateTime<chrono::Utc>>,
pub model_architecture: Option<String>,
pub dataset_id: Option<String>,
pub config: HashMap<String, String>,
pub metrics: SessionMetrics,
pub code_history: Vec<CodeCell>,
pub tags: Vec<String>,
pub notes: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeCell {
pub execution_order: u32,
pub source: String,
pub output: Option<String>,
pub timestamp: chrono::DateTime<chrono::Utc>,
pub duration_ms: Option<u64>,
}
impl EntrenarSession {
pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
Self {
id: id.into(),
name: name.into(),
user: None,
created_at: chrono::Utc::now(),
ended_at: None,
model_architecture: None,
dataset_id: None,
config: HashMap::new(),
metrics: SessionMetrics::new(),
code_history: Vec::new(),
tags: Vec::new(),
notes: None,
}
}
pub fn with_user(mut self, user: impl Into<String>) -> Self {
self.user = Some(user.into());
self
}
pub fn with_architecture(mut self, arch: impl Into<String>) -> Self {
self.model_architecture = Some(arch.into());
self
}
pub fn with_dataset(mut self, dataset: impl Into<String>) -> Self {
self.dataset_id = Some(dataset.into());
self
}
pub fn with_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.config.insert(key.into(), value.into());
self
}
pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
self.tags.push(tag.into());
self
}
pub fn with_notes(mut self, notes: impl Into<String>) -> Self {
self.notes = Some(notes.into());
self
}
pub fn add_code_cell(&mut self, cell: CodeCell) {
self.code_history.push(cell);
}
pub fn end(&mut self) {
self.ended_at = Some(chrono::Utc::now());
}
pub fn duration(&self) -> Option<chrono::Duration> {
self.ended_at.map(|end| end - self.created_at)
}
pub fn has_training_data(&self) -> bool {
!self.metrics.is_empty()
}
pub fn export_json(&self) -> Result<serde_json::Value, serde_json::Error> {
let export = SessionExport {
session_id: self.id.clone(),
name: self.name.clone(),
user: self.user.clone(),
created_at: self.created_at.to_rfc3339(),
ended_at: self.ended_at.map(|t| t.to_rfc3339()),
duration_seconds: self.duration().map(|d| d.num_seconds()),
model_architecture: self.model_architecture.clone(),
dataset_id: self.dataset_id.clone(),
config: self.config.clone(),
metrics: MetricsExportSummary {
total_steps: self.metrics.total_steps(),
final_loss: self.metrics.final_loss(),
best_loss: self.metrics.best_loss(),
final_accuracy: self.metrics.final_accuracy(),
best_accuracy: self.metrics.best_accuracy(),
loss_history: self.metrics.loss_history.clone(),
accuracy_history: self.metrics.accuracy_history.clone(),
custom_metrics: self.metrics.custom.clone(),
},
code_cells_count: self.code_history.len(),
tags: self.tags.clone(),
notes: self.notes.clone(),
};
serde_json::to_value(export)
}
pub fn export_json_string(&self) -> Result<String, serde_json::Error> {
let value = self.export_json()?;
serde_json::to_string_pretty(&value)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionExport {
pub session_id: String,
pub name: String,
pub user: Option<String>,
pub created_at: String,
pub ended_at: Option<String>,
pub duration_seconds: Option<i64>,
pub model_architecture: Option<String>,
pub dataset_id: Option<String>,
pub config: HashMap<String, String>,
pub metrics: MetricsExportSummary,
pub code_cells_count: usize,
pub tags: Vec<String>,
pub notes: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MetricsExportSummary {
pub total_steps: usize,
pub final_loss: Option<f64>,
pub best_loss: Option<f64>,
pub final_accuracy: Option<f64>,
pub best_accuracy: Option<f64>,
pub loss_history: Vec<f64>,
pub accuracy_history: Vec<f64>,
pub custom_metrics: HashMap<String, Vec<f64>>,
}