forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
"""Test LangChain+ Client Utils."""
|
|
|
|
from typing import List
|
|
|
|
from langchain.client.utils import parse_chat_messages
|
|
from langchain.schema import (
|
|
AIMessage,
|
|
BaseMessage,
|
|
ChatMessage,
|
|
HumanMessage,
|
|
SystemMessage,
|
|
)
|
|
|
|
|
|
def test_parse_chat_messages() -> None:
|
|
"""Test that chat messages are parsed correctly."""
|
|
input_text = (
|
|
"Human: I am human roar\nAI: I am AI beep boop\nSystem: I am a system message"
|
|
)
|
|
expected = [
|
|
HumanMessage(content="I am human roar"),
|
|
AIMessage(content="I am AI beep boop"),
|
|
SystemMessage(content="I am a system message"),
|
|
]
|
|
assert parse_chat_messages(input_text) == expected
|
|
|
|
|
|
def test_parse_chat_messages_empty_input() -> None:
|
|
"""Test that an empty input string returns an empty list."""
|
|
input_text = ""
|
|
expected: List[BaseMessage] = []
|
|
assert parse_chat_messages(input_text) == expected
|
|
|
|
|
|
def test_parse_chat_messages_multiline_messages() -> None:
|
|
"""Test that multiline messages are parsed correctly."""
|
|
input_text = (
|
|
"Human: I am a human\nand I roar\nAI: I am an AI\nand I"
|
|
" beep boop\nSystem: I am a system\nand a message"
|
|
)
|
|
expected = [
|
|
HumanMessage(content="I am a human\nand I roar"),
|
|
AIMessage(content="I am an AI\nand I beep boop"),
|
|
SystemMessage(content="I am a system\nand a message"),
|
|
]
|
|
assert parse_chat_messages(input_text) == expected
|
|
|
|
|
|
def test_parse_chat_messages_custom_roles() -> None:
|
|
"""Test that custom roles are parsed correctly."""
|
|
input_text = "Client: I need help\nAgent: I'm here to help\nClient: Thank you"
|
|
expected = [
|
|
ChatMessage(role="Client", content="I need help"),
|
|
ChatMessage(role="Agent", content="I'm here to help"),
|
|
ChatMessage(role="Client", content="Thank you"),
|
|
]
|
|
assert parse_chat_messages(input_text, roles=["Client", "Agent"]) == expected
|
|
|
|
|
|
def test_parse_chat_messages_embedded_roles() -> None:
|
|
"""Test that messages with embedded role references are parsed correctly."""
|
|
input_text = (
|
|
"Human: Oh ai what if you said AI: foo bar?"
|
|
"\nAI: Well, that would be interesting!"
|
|
)
|
|
expected = [
|
|
HumanMessage(content="Oh ai what if you said AI: foo bar?"),
|
|
AIMessage(content="Well, that would be interesting!"),
|
|
]
|
|
assert parse_chat_messages(input_text) == expected
|