refactor: Message

This commit is contained in:
莘权 马 2023-10-31 15:23:37 +08:00
parent d8adba99d4
commit 5e8ada5cff
2 changed files with 92 additions and 36 deletions

View file

@ -11,26 +11,30 @@ from metagpt.schema import AIMessage, Message, RawMessage, SystemMessage, UserMe
def test_message():
msg = Message(role='User', content='WTF')
assert msg.to_dict()['role'] == 'User'
assert 'User' in str(msg)
msg = Message(role="User", content="WTF")
assert msg.to_dict()["role"] == "User"
assert "User" in str(msg)
def test_all_messages():
test_content = 'test_message'
test_content = "test_message"
msgs = [
UserMessage(test_content),
SystemMessage(test_content),
AIMessage(test_content),
Message(test_content, role='QA')
Message(test_content, role="QA"),
]
for msg in msgs:
assert msg.content == test_content
def test_raw_message():
msg = RawMessage(role='user', content='raw')
assert msg['role'] == 'user'
assert msg['content'] == 'raw'
msg = RawMessage(role="user", content="raw")
assert msg["role"] == "user"
assert msg["content"] == "raw"
with pytest.raises(KeyError):
assert msg['1'] == 1, "KeyError: '1'"
assert msg["1"] == 1, "KeyError: '1'"
if __name__ == "__main__":
pytest.main([__file__, "-s"])