Spaces:
Running
Running
File size: 2,216 Bytes
cfd3735 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
"""Test formatting functionality."""
import unittest
from langchain.schema import (
AIMessage,
HumanMessage,
SystemMessage,
get_buffer_string,
messages_from_dict,
messages_to_dict,
)
class TestGetBufferString(unittest.TestCase):
human_msg: HumanMessage = HumanMessage(content="human")
ai_msg: AIMessage = AIMessage(content="ai")
sys_msg: SystemMessage = SystemMessage(content="sys")
def test_empty_input(self) -> None:
self.assertEqual(get_buffer_string([]), "")
def test_valid_single_message(self) -> None:
expected_output = f"Human: {self.human_msg.content}"
self.assertEqual(
get_buffer_string([self.human_msg]),
expected_output,
)
def test_custom_human_prefix(self) -> None:
prefix = "H"
expected_output = f"{prefix}: {self.human_msg.content}"
self.assertEqual(
get_buffer_string([self.human_msg], human_prefix="H"),
expected_output,
)
def test_custom_ai_prefix(self) -> None:
prefix = "A"
expected_output = f"{prefix}: {self.ai_msg.content}"
self.assertEqual(
get_buffer_string([self.ai_msg], ai_prefix="A"),
expected_output,
)
def test_multiple_msg(self) -> None:
msgs = [self.human_msg, self.ai_msg, self.sys_msg]
expected_output = "\n".join(
[
f"Human: {self.human_msg.content}",
f"AI: {self.ai_msg.content}",
f"System: {self.sys_msg.content}",
]
)
self.assertEqual(
get_buffer_string(msgs),
expected_output,
)
class TestMessageDictConversion(unittest.TestCase):
human_msg: HumanMessage = HumanMessage(
content="human", additional_kwargs={"key": "value"}
)
ai_msg: AIMessage = AIMessage(content="ai")
sys_msg: SystemMessage = SystemMessage(content="sys")
def test_multiple_msg(self) -> None:
msgs = [
self.human_msg,
self.ai_msg,
self.sys_msg,
]
self.assertEqual(
messages_from_dict(messages_to_dict(msgs)),
msgs,
)
|