Skip to content

Commit

Permalink
Only set chat_history if there is None
Browse files Browse the repository at this point in the history
This allows for the chat_history not to be blown away by a default instance.
  • Loading branch information
boydjohnson committed Dec 15, 2024
1 parent b20a998 commit 7ec2d7b
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 8 deletions.
13 changes: 6 additions & 7 deletions src/generation/chat/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -207,15 +207,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);
&mut self.messages_history.clone().unwrap()
}
};
// Clone the current chat messages to avoid borrowing issues
// And not to add message to the history if the request fails
Expand Down
41 changes: 40 additions & 1 deletion tests/chat_history_management.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
use ollama_rs::Ollama;
use ollama_rs::{
generation::chat::{request::ChatMessageRequest, ChatMessage, MessageRole},
Ollama,
};

#[test]
fn test_chat_history_saved_as_should() {
Expand Down Expand Up @@ -111,3 +114,39 @@ fn test_chat_history_freed_if_limit_exceeded() {
assert!(last.is_some());
assert_eq!(last.unwrap().content, "Hi again".to_string());
}

#[tokio::test]
async fn test_chat_history_accumulated() {
let mut ollama = Ollama::new_default_with_history(30);
let chat_id = "default";

assert!(ollama
.send_chat_messages_with_history(
ChatMessageRequest::new(
"granite-code:3b".into(),
vec![ChatMessage::new(
MessageRole::User,
"Why is the sky blue?".into(),
)],
),
chat_id,
)
.await
.is_ok());

assert!(ollama
.send_chat_messages_with_history(
ChatMessageRequest::new(
"granite-code:3b".into(),
vec![ChatMessage::new(
MessageRole::User,
"But, why is the sky blue?".into()
)]
),
chat_id
)
.await
.is_ok());

assert_eq!(ollama.get_messages_history(chat_id).unwrap().len(), 4)
}

0 comments on commit 7ec2d7b

Please sign in to comment.