Fix panic when removing a message in the history with 0 value

[shigedangao]

Description

It appears that when calling the send_chat_messages_with_history_stream method. Theadd_message method will try to remove a message by finding an index. Which by default is set to 0.

ollama-rs/src/history.rs

Line 42 in 1900774

messages.remove(index_to_remove);

For some unknown reason the index_to_remove variable can be equal to 0 and the fact that the messages vec can also be empty this means that the operation will panic

This PR aims to just to avoid panicking and only remove a message when the id is present

[starccy]

Hi, thanks for your work! After a closer inspection, however, I am afraid that this PR did not fix the root issue.

The root problem I think is every time it tries to get the stored history, the history will be reassigned with a default value.

ollama-rs/src/generation/chat/mod.rs

Lines 198 to 207 in 1900774

	fn get_chat_messages_by_id(&mut self, history_id: impl ToString) -> Vec<ChatMessage> {
	let mut binding = {
	let new_history =
	std::sync::Arc::new(std::sync::RwLock::new(MessagesHistory::default()));
	self.messages_history = Some(new_history);
	self.messages_history.clone().unwrap()
	};
	let chat_history = match self.messages_history.as_mut() {
	Some(history) => history,
	None => &mut binding,

Because of this, it not only caused the panic problem but the historical messages were not stored at all.

How about moving the binding statement into the match block?

    fn get_chat_messages_by_id(&mut self, history_id: impl ToString) -> Vec<ChatMessage> {
        let chat_history = match self.messages_history.as_mut() {
            Some(history) => history,
            None => {
                let new_history =
                    std::sync::Arc::new(std::sync::RwLock::new(MessagesHistory::default()));
                self.messages_history = Some(new_history.clone());
                self.messages_history.as_mut().unwrap()
            },
        };

Besides, give the MessagesHistory a default limits value to prevent array index overflow.

#[derive(Debug, Clone)]
pub struct MessagesHistory {
    pub(crate) messages_by_id: HashMap<String, Vec<ChatMessage>>,
    pub(crate) messages_number_limit: u16,
}

impl Default for MessagesHistory {
    fn default() -> Self {
        Self {
            messages_by_id: HashMap::new(),
            messages_number_limit: 30,
        }
    }
}

[starccy]

Here's the diff, hoping it's helpful.

diff --git a/src/generation/chat/mod.rs b/src/generation/chat/mod.rs
index 7fdceb0..dbe8b86 100644
--- a/src/generation/chat/mod.rs
+++ b/src/generation/chat/mod.rs
@@ -196,15 +196,14 @@ impl Ollama {
     /// Without impact for existing history
     /// Used to prepare history for request
     fn get_chat_messages_by_id(&mut self, history_id: impl ToString) -> Vec<ChatMessage> {
-        let mut binding = {
-            let new_history =
-                std::sync::Arc::new(std::sync::RwLock::new(MessagesHistory::default()));
-            self.messages_history = Some(new_history);
-            self.messages_history.clone().unwrap()
-        };
         let chat_history = match self.messages_history.as_mut() {
             Some(history) => history,
-            None => &mut binding,
+            None => {
+                let new_history =
+                    std::sync::Arc::new(std::sync::RwLock::new(MessagesHistory::default()));
+                self.messages_history = Some(new_history.clone());
+                self.messages_history.as_mut().unwrap()
+            },
         };
         // Clone the current chat messages to avoid borrowing issues
         // And not to add message to the history if the request fails
diff --git a/src/history.rs b/src/history.rs
index 2fe3eea..a90924c 100644
--- a/src/history.rs
+++ b/src/history.rs
@@ -5,12 +5,21 @@ use crate::{
     Ollama,
 };
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 pub struct MessagesHistory {
     pub(crate) messages_by_id: HashMap<String, Vec<ChatMessage>>,
     pub(crate) messages_number_limit: u16,
 }
 
+impl Default for MessagesHistory {
+    fn default() -> Self {
+        Self {
+            messages_by_id: HashMap::new(),
+            messages_number_limit: 30,
+        }
+    }
+}
+
 pub type WrappedMessageHistory = std::sync::Arc<std::sync::RwLock<MessagesHistory>>;
 
 /// Store for messages history

[bredmond1019]

Hey all. I just found the same issue and the same solution on local, then saw you had this change already. I'm not sure it
s necessary to set the default always to 30 like the diff above, but the rest checks out on my end.

In the send_chat_messages_with_history_stream function, the get_chat_messages_by_id is always called. Then the binding variable will always set the messages_history back to a new one, so whatever we set in the beginning won't matter.

[shigedangao]

Hey thanks for your feedback. I indeed haven't dig quite far into the code and it's good that you guys did as it'd have not fix the root cause issue. I'm not sure however if it's necessary to add a Default value, as when creating the Ollama instance the history limit would have been set e.g: with new_with_history.

Let me know if the fix works on your end @bredmond1019 @starccy

[shigedangao]

Actually a similar PR seems to fix the issue. See here #60

[starccy]

Oh yes, you are right, the default value is not necessary. And I realize there is already a PR to solve this problem indeed.

[shigedangao]

@starccy closing this one as the previous PR will fix our issue