diff --git a/src/history.rs b/src/history.rs index 2fe3eea..0d79832 100644 --- a/src/history.rs +++ b/src/history.rs @@ -34,12 +34,11 @@ impl MessagesHistory { // Replacing the oldest message if the limit is reached // The oldest message is the first one, unless it's a system message if messages.len() >= self.messages_number_limit as usize { - let index_to_remove = messages - .first() - .map(|m| if m.role == MessageRole::System { 1 } else { 0 }) - .unwrap_or(0); - - messages.remove(index_to_remove); + if let Some(index_to_remove) = + messages.iter().position(|m| m.role != MessageRole::System) + { + messages.remove(index_to_remove); + } } if message.role == MessageRole::System { diff --git a/tests/chat_history_management.rs b/tests/chat_history_management.rs index cd98d03..a1935bd 100644 --- a/tests/chat_history_management.rs +++ b/tests/chat_history_management.rs @@ -88,3 +88,26 @@ fn clear_chat_history_for_all() { assert!(ollama.get_messages_history(first_chat_id).is_none()); assert!(ollama.get_messages_history(another_chat_id).is_none()); } + +#[test] +fn test_chat_history_freed_if_limit_exceeded() { + let mut ollama = Ollama::new_default_with_history(3); + let chat_id = "default"; + + ollama.add_user_response(chat_id, "Hello"); + ollama.add_assistant_response(chat_id, "Hi"); + + ollama.add_user_response(chat_id, "Tell me 'hi' again"); + ollama.add_assistant_response(chat_id, "Hi again"); + + ollama.add_user_response(chat_id, "Tell me 'hi' again"); + ollama.add_assistant_response(chat_id, "Hi again"); + + let history = ollama.get_messages_history(chat_id).unwrap(); + + assert_eq!(history.len(), 3); + + let last = history.last(); + assert!(last.is_some()); + assert_eq!(last.unwrap().content, "Hi again".to_string()); +}