use std::collections::VecDeque;
#[derive(Debug, Clone)]
pub struct Message {
pub id: String,
pub role: MessageRole,
pub content: String,
pub timestamp: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub enum MessageRole {
User,
Assistant,
System,
Tool,
}
impl Message {
pub fn user(content: &str) -> Self {
Self::new("user", content)
}
pub fn assistant(content: &str) -> Self {
Self::new("assistant", content)
}
pub fn system(content: &str) -> Self {
Self::new("system", content)
}
pub fn tool(content: &str) -> Self {
Self::new("tool", content)
}
fn new(role: &str, content: &str) -> Self {
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_millis() as u64;
Self {
id: uuid::Uuid::new_v4().to_string(),
role: match role {
"user" => MessageRole::User,
"assistant" => MessageRole::Assistant,
"system" => MessageRole::System,
"tool" => MessageRole::Tool,
_ => MessageRole::User,
},
content: content.to_string(),
timestamp,
}
}
}
pub struct MessageHistory {
messages: VecDeque<Message>,
max_size: usize,
}
impl MessageHistory {
pub fn new(max_size: usize) -> Self {
Self {
messages: VecDeque::new(),
max_size,
}
}
pub fn add(&mut self, message: Message) {
if self.messages.len() >= self.max_size {
self.messages.pop_front();
}
self.messages.push_back(message);
}
pub fn get_messages(&self) -> Vec<&Message> {
self.messages.iter().collect()
}
pub fn clear(&mut self) {
self.messages.clear();
}
}