mirror of
https://github.com/hwchase17/langchain
synced 2024-10-29 17:07:25 +00:00
3bfe7cf467
should be no functional changes also keep __init__ exposing a lot for backwards compat --------- Co-authored-by: Dev 2049 <dev.dev2049@gmail.com> Co-authored-by: Bagatur <baskaryan@gmail.com>
78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
"""Test formatting functionality."""
|
|
|
|
import unittest
|
|
|
|
from langchain.schema.messages 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,
|
|
)
|