"""Test Alibaba Tongyi Chat Model.""" from langchain_core.callbacks import CallbackManager from langchain_core.messages import AIMessage, BaseMessage, HumanMessage from langchain_core.outputs import ChatGeneration, LLMResult from langchain_core.pydantic_v1 import SecretStr from pytest import CaptureFixture from langchain_community.chat_models.tongyi import ChatTongyi from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler def test_api_key_is_string() -> None: llm = ChatTongyi(dashscope_api_key="secret-api-key") assert isinstance(llm.dashscope_api_key, SecretStr) def test_api_key_masked_when_passed_via_constructor( capsys: CaptureFixture, ) -> None: llm = ChatTongyi(dashscope_api_key="secret-api-key") print(llm.dashscope_api_key, end="") captured = capsys.readouterr() assert captured.out == "**********" def test_default_call() -> None: """Test default model call.""" chat = ChatTongyi() response = chat(messages=[HumanMessage(content="Hello")]) assert isinstance(response, BaseMessage) assert isinstance(response.content, str) def test_model() -> None: """Test model kwarg works.""" chat = ChatTongyi(model="qwen-plus") response = chat(messages=[HumanMessage(content="Hello")]) assert isinstance(response, BaseMessage) assert isinstance(response.content, str) def test_multiple_history() -> None: """Tests multiple history works.""" chat = ChatTongyi() response = chat( messages=[ HumanMessage(content="Hello."), AIMessage(content="Hello!"), HumanMessage(content="How are you doing?"), ] ) assert isinstance(response, BaseMessage) assert isinstance(response.content, str) def test_stream() -> None: """Test that stream works.""" chat = ChatTongyi(streaming=True) callback_handler = FakeCallbackHandler() callback_manager = CallbackManager([callback_handler]) response = chat( messages=[ HumanMessage(content="Hello."), AIMessage(content="Hello!"), HumanMessage(content="Who are you?"), ], stream=True, callbacks=callback_manager, ) assert callback_handler.llm_streams > 0 assert isinstance(response.content, str) def test_multiple_messages() -> None: """Tests multiple messages works.""" chat = ChatTongyi() message = HumanMessage(content="Hi, how are you.") response = chat.generate([[message], [message]]) assert isinstance(response, LLMResult) assert len(response.generations) == 2 for generations in response.generations: assert len(generations) == 1 for generation in generations: assert isinstance(generation, ChatGeneration) assert isinstance(generation.text, str) assert generation.text == generation.message.content