langchain/libs/community/tests/integration_tests/chat_models/test_sparkllm.py
Eugene Yurtsev 25fbe356b4
community[patch]: upgrade to recent version of mypy (#21616)
This PR upgrades community to a recent version of mypy. It inserts type:
ignore on all existing failures.
2024-05-13 14:55:07 -04:00

46 lines
1.5 KiB
Python

from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage
from langchain_community.chat_models.sparkllm import ChatSparkLLM
def test_initialization() -> None:
"""Test chat model initialization."""
for model in [
ChatSparkLLM(timeout=30),
ChatSparkLLM(request_timeout=30), # type: ignore[call-arg]
]:
assert model.request_timeout == 30
def test_chat_spark_llm() -> None:
chat = ChatSparkLLM() # type: ignore[call-arg]
message = HumanMessage(content="Hello")
response = chat.invoke([message])
assert isinstance(response, AIMessage)
assert isinstance(response.content, str)
def test_chat_spark_llm_streaming() -> None:
chat = ChatSparkLLM(streaming=True) # type: ignore[call-arg]
for chunk in chat.stream("Hello!"):
assert isinstance(chunk, AIMessageChunk)
assert isinstance(chunk.content, str)
def test_chat_spark_llm_with_domain() -> None:
chat = ChatSparkLLM(spark_llm_domain="generalv3") # type: ignore[call-arg]
message = HumanMessage(content="Hello")
response = chat.invoke([message])
print(response) # noqa: T201
assert isinstance(response, AIMessage)
assert isinstance(response.content, str)
def test_chat_spark_llm_with_temperature() -> None:
chat = ChatSparkLLM(temperature=0.9, top_k=2) # type: ignore[call-arg]
message = HumanMessage(content="Hello")
response = chat.invoke([message])
print(response) # noqa: T201
assert isinstance(response, AIMessage)
assert isinstance(response.content, str)