use chrono::Utc;
use uuid::Uuid;
use super::reader::{read_storage_file, write_storage_file};
use crate::error::StorageError;
use crate::types::StorageSession;
pub fn create_new_session() -> Result<String, StorageError> {
let mut data = read_storage_file()?;
let session_id = Uuid::new_v4().to_string();
let new_session = StorageSession {
session_id: session_id.clone(),
status: 0,
create_at: Utc::now().to_rfc3339(),
entries: Vec::new(),
compacted_entries: Vec::new(),
};
data.sessions.push(new_session);
write_storage_file(&data)?;
Ok(session_id)
}
pub fn write_entry_to_file(
session_id: &str,
role: &str,
content: &str,
tools: Option<&serde_json::Value>,
token_consumption: Option<u32>,
create_at: Option<&str>,
) -> Result<(), StorageError> {
let mut data = read_storage_file()?;
let session = data
.sessions
.iter_mut()
.find(|s| s.session_id == session_id);
match session {
Some(s) => {
let entry_id = Uuid::new_v4().to_string();
let new_entry = crate::types::Entry {
entry_id,
session_id: session_id.to_string(),
content: content.to_string(),
role: role.to_string(),
token_consumption,
status: 0,
tools: tools.cloned(),
create_at: create_at
.map(|s| s.to_string())
.unwrap_or_else(|| Utc::now().to_rfc3339()),
is_compaction: 1,
};
s.entries.push(new_entry);
}
None => {
let entry_id = Uuid::new_v4().to_string();
let new_entry = crate::types::Entry {
entry_id,
session_id: session_id.to_string(),
content: content.to_string(),
role: role.to_string(),
token_consumption,
status: 0,
tools: tools.cloned(),
create_at: create_at
.map(|s| s.to_string())
.unwrap_or_else(|| Utc::now().to_rfc3339()),
is_compaction: 1,
};
let new_session = StorageSession {
session_id: session_id.to_string(),
status: 0,
create_at: Utc::now().to_rfc3339(),
entries: vec![new_entry],
compacted_entries: Vec::new(),
};
data.sessions.push(new_session);
}
}
write_storage_file(&data)?;
Ok(())
}
pub fn write_compacted_entry_to_file(
session_id: &str,
summary: &str,
trigger_entry_id: &str,
create_at: Option<&str>,
) -> Result<(), StorageError> {
let mut data = read_storage_file()?;
let session = data
.sessions
.iter_mut()
.find(|s| s.session_id == session_id);
let entry_id = Uuid::new_v4().to_string();
let new_compacted = crate::types::CompactedEntry {
entry_id,
session_id: session_id.to_string(),
trigger_entry_id: trigger_entry_id.to_string(),
summary: summary.to_string(),
create_at: create_at
.map(|s| s.to_string())
.unwrap_or_else(|| Utc::now().to_rfc3339()),
status: 0,
};
match session {
Some(s) => {
s.compacted_entries.push(new_compacted);
}
None => {
let new_session = StorageSession {
session_id: session_id.to_string(),
status: 0,
create_at: Utc::now().to_rfc3339(),
entries: Vec::new(),
compacted_entries: vec![new_compacted],
};
data.sessions.push(new_session);
}
}
write_storage_file(&data)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_create_new_session() {
let session_id = create_new_session().unwrap();
assert!(!session_id.is_empty());
assert!(session_id.contains('-'));
}
}