langchain/libs/community/tests/unit_tests/chat_models/test_yuan2.py
Austin Burdette f355a98bb6
community:yuan2[patch]: standardize init args (#21462)
updated stop and request_timeout so they aliased to stop_sequences, and
timeout respectively. Added test that both continue to set the same
underlying attributes.

Related to
[20085](https://github.com/langchain-ai/langchain/issues/20085)

Co-authored-by: ccurme <chester.curme@gmail.com>
2024-08-23 17:56:19 +00:00

81 lines
2.5 KiB
Python

"""Test ChatYuan2 wrapper."""
import pytest
from langchain_core.messages import (
AIMessage,
HumanMessage,
SystemMessage,
)
from langchain_community.chat_models.yuan2 import (
ChatYuan2,
_convert_dict_to_message,
_convert_message_to_dict,
)
@pytest.mark.requires("openai")
def test_yuan2_model_param() -> None:
chat = ChatYuan2(model="foo") # type: ignore[call-arg]
assert chat.model_name == "foo"
chat = ChatYuan2(model_name="foo") # type: ignore[call-arg]
assert chat.model_name == "foo"
@pytest.mark.requires("openai")
def test_yuan2_timeout_param() -> None:
chat = ChatYuan2(request_timeout=5) # type: ignore[call-arg]
assert chat.request_timeout == 5
chat = ChatYuan2(timeout=10) # type: ignore[call-arg]
assert chat.request_timeout == 10
@pytest.mark.requires("openai")
def test_yuan2_stop_sequences_param() -> None:
chat = ChatYuan2(stop=["<eod>"]) # type: ignore[call-arg]
assert chat.stop == ["<eod>"]
chat = ChatYuan2(stop_sequences=["<eod>"]) # type: ignore[call-arg]
assert chat.stop == ["<eod>"]
def test__convert_message_to_dict_human() -> None:
message = HumanMessage(content="foo")
result = _convert_message_to_dict(message)
expected_output = {"role": "user", "content": "foo"}
assert result == expected_output
def test__convert_message_to_dict_ai() -> None:
message = AIMessage(content="foo")
result = _convert_message_to_dict(message)
expected_output = {"role": "assistant", "content": "foo"}
assert result == expected_output
def test__convert_message_to_dict_system() -> None:
message = SystemMessage(content="foo")
result = _convert_message_to_dict(message)
expected_output = {"role": "system", "content": "foo"}
assert result == expected_output
def test__convert_dict_to_message_human() -> None:
message = {"role": "user", "content": "hello"}
result = _convert_dict_to_message(message)
expected_output = HumanMessage(content="hello")
assert result == expected_output
def test__convert_dict_to_message_ai() -> None:
message = {"role": "assistant", "content": "hello"}
result = _convert_dict_to_message(message)
expected_output = AIMessage(content="hello")
assert result == expected_output
def test__convert_dict_to_message_system() -> None:
message = {"role": "system", "content": "hello"}
result = _convert_dict_to_message(message)
expected_output = SystemMessage(content="hello")
assert result == expected_output