[Improvement] Use SQLite for chat memory (#910)
Co-authored-by: Deven Patel <deven298@yahoo.com>
This commit is contained in:
67
tests/memory/test_chat_memory.py
Normal file
67
tests/memory/test_chat_memory.py
Normal file
@@ -0,0 +1,67 @@
|
||||
import pytest
|
||||
|
||||
from embedchain.memory.base import ECChatMemory
|
||||
from embedchain.memory.message import ChatMessage
|
||||
|
||||
|
||||
# Fixture for creating an instance of ECChatMemory
|
||||
@pytest.fixture
|
||||
def chat_memory_instance():
|
||||
return ECChatMemory()
|
||||
|
||||
|
||||
def test_add_chat_memory(chat_memory_instance):
|
||||
app_id = "test_app"
|
||||
human_message = "Hello, how are you?"
|
||||
ai_message = "I'm fine, thank you!"
|
||||
|
||||
chat_message = ChatMessage()
|
||||
chat_message.add_user_message(human_message)
|
||||
chat_message.add_ai_message(ai_message)
|
||||
|
||||
chat_memory_instance.add(app_id, chat_message)
|
||||
|
||||
assert chat_memory_instance.count_history_messages(app_id) == 1
|
||||
chat_memory_instance.delete_chat_history(app_id)
|
||||
|
||||
|
||||
def test_get_recent_memories(chat_memory_instance):
|
||||
app_id = "test_app"
|
||||
|
||||
for i in range(1, 7):
|
||||
human_message = f"Question {i}"
|
||||
ai_message = f"Answer {i}"
|
||||
|
||||
chat_message = ChatMessage()
|
||||
chat_message.add_user_message(human_message)
|
||||
chat_message.add_ai_message(ai_message)
|
||||
|
||||
chat_memory_instance.add(app_id, chat_message)
|
||||
|
||||
recent_memories = chat_memory_instance.get_recent_memories(app_id, num_rounds=5)
|
||||
|
||||
assert len(recent_memories) == 5
|
||||
|
||||
|
||||
def test_delete_chat_history(chat_memory_instance):
|
||||
app_id = "test_app"
|
||||
|
||||
for i in range(1, 6):
|
||||
human_message = f"Question {i}"
|
||||
ai_message = f"Answer {i}"
|
||||
|
||||
chat_message = ChatMessage()
|
||||
chat_message.add_user_message(human_message)
|
||||
chat_message.add_ai_message(ai_message)
|
||||
|
||||
chat_memory_instance.add(app_id, chat_message)
|
||||
|
||||
chat_memory_instance.delete_chat_history(app_id)
|
||||
|
||||
assert chat_memory_instance.count_history_messages(app_id) == 0
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def close_connection(chat_memory_instance):
|
||||
yield
|
||||
chat_memory_instance.close_connection()
|
||||
37
tests/memory/test_memory_messages.py
Normal file
37
tests/memory/test_memory_messages.py
Normal file
@@ -0,0 +1,37 @@
|
||||
from embedchain.memory.message import BaseMessage, ChatMessage
|
||||
|
||||
|
||||
def test_ec_base_message():
|
||||
content = "Hello, how are you?"
|
||||
creator = "human"
|
||||
metadata = {"key": "value"}
|
||||
|
||||
message = BaseMessage(content=content, creator=creator, metadata=metadata)
|
||||
|
||||
assert message.content == content
|
||||
assert message.creator == creator
|
||||
assert message.metadata == metadata
|
||||
assert message.type is None
|
||||
assert message.is_lc_serializable() is True
|
||||
assert str(message) == f"{creator}: {content}"
|
||||
|
||||
|
||||
def test_ec_base_chat_message():
|
||||
human_message_content = "Hello, how are you?"
|
||||
ai_message_content = "I'm fine, thank you!"
|
||||
human_metadata = {"user": "John"}
|
||||
ai_metadata = {"response_time": 0.5}
|
||||
|
||||
chat_message = ChatMessage()
|
||||
chat_message.add_user_message(human_message_content, metadata=human_metadata)
|
||||
chat_message.add_ai_message(ai_message_content, metadata=ai_metadata)
|
||||
|
||||
assert chat_message.human_message.content == human_message_content
|
||||
assert chat_message.human_message.creator == "human"
|
||||
assert chat_message.human_message.metadata == human_metadata
|
||||
|
||||
assert chat_message.ai_message.content == ai_message_content
|
||||
assert chat_message.ai_message.creator == "ai"
|
||||
assert chat_message.ai_message.metadata == ai_metadata
|
||||
|
||||
assert str(chat_message) == f"human: {human_message_content} | ai: {ai_message_content}"
|
||||
Reference in New Issue
Block a user