use rexis_llm::{ChatMessage, MessageRole};
#[derive(Debug, Clone)]
pub struct ConversationMemory {
messages: Vec<ChatMessage>,
max_length: usize,
}
impl ConversationMemory {
pub fn new() -> Self {
Self {
messages: Vec::new(),
max_length: 50,
}
}
pub fn with_max_length(max_length: usize) -> Self {
Self {
messages: Vec::new(),
max_length,
}
}
pub fn add_message(&mut self, message: ChatMessage) {
self.messages.push(message);
if self.messages.len() > self.max_length {
let system_msg = self.messages.first().cloned();
self.messages
.drain(1..self.messages.len() - self.max_length + 1);
if let Some(sys_msg) = system_msg {
if matches!(sys_msg.role, MessageRole::System) {
if !matches!(self.messages.first().unwrap().role, MessageRole::System) {
self.messages.insert(0, sys_msg);
}
}
}
}
}
pub fn get_messages(&self) -> &[ChatMessage] {
&self.messages
}
pub fn to_messages(&self) -> Vec<ChatMessage> {
self.messages.clone()
}
pub fn clear(&mut self) {
let system_msg = self
.messages
.iter()
.find(|m| matches!(m.role, MessageRole::System))
.cloned();
self.messages.clear();
if let Some(msg) = system_msg {
self.messages.push(msg);
}
}
pub fn len(&self) -> usize {
self.messages.len()
}
pub fn is_empty(&self) -> bool {
self.messages.is_empty()
}
}
impl Default for ConversationMemory {
fn default() -> Self {
Self::new()
}
}