Chat buffer should return last messages, not first (#6979)

run-llama · Jul 20, 2023 · 991c0eb · 991c0eb
1 parent be0ded3
commit 991c0eb
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,10 @@
 # ChangeLog
 
+## Unreleased
+
+### Bug Fixes / Nits
+- Fix returned order of messages in large chat memory (#6979) 
+
 ## [v0.7.11] - 2023-07-19
 
 ### New Features

diff --git a/llama_index/memory/chat_memory_buffer.py b/llama_index/memory/chat_memory_buffer.py
@@ -56,7 +56,7 @@ def get(self) -> List[ChatMessage]:
         if token_count > self.token_limit:
             return []
 
-        return self.chat_history[:message_count]
+        return self.chat_history[-message_count:]
 
     def get_all(self) -> List[ChatMessage]:
         """Get all chat history."""

diff --git a/tests/memory/test_chat_memory_buffer.py b/tests/memory/test_chat_memory_buffer.py
@@ -27,7 +27,7 @@ def test_set() -> None:
 def test_max_tokens() -> None:
     memory = ChatMemoryBuffer.from_defaults(chat_history=[CHAT_MESSAGE], token_limit=5)
 
-    memory.put(ChatMessage(role=MessageRole.USER, content="test message2"))
+    memory.put(CHAT_MESSAGE)
     assert len(memory.get()) == 2
 
     # do we limit properly
@@ -39,4 +39,6 @@ def test_max_tokens() -> None:
     assert len(memory.get_all()) == 4
 
     # does get return in the correct order?
+    memory.put(ChatMessage(role=MessageRole.USER, content="test message2"))
     assert memory.get()[-1].content == "test message2"
+    assert len(memory.get()) == 2