langchain/libs/community/tests/integration_tests/llms/test_qianfan_endpoint.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

44 lines
1.4 KiB
Python

"""Test Baidu Qianfan LLM Endpoint."""
from typing import Generator
from langchain_core.outputs import LLMResult
from langchain_community.llms.baidu_qianfan_endpoint import QianfanLLMEndpoint
def test_call() -> None:
"""Test valid call to qianfan."""
llm = QianfanLLMEndpoint() # type: ignore[call-arg]
output = llm.invoke("write a joke")
assert isinstance(output, str)
def test_generate() -> None:
"""Test valid call to qianfan."""
llm = QianfanLLMEndpoint() # type: ignore[call-arg]
output = llm.generate(["write a joke"])
assert isinstance(output, LLMResult)
assert isinstance(output.generations, list)
def test_generate_stream() -> None:
"""Test valid call to qianfan."""
llm = QianfanLLMEndpoint() # type: ignore[call-arg]
output = llm.stream("write a joke")
assert isinstance(output, Generator)
async def test_qianfan_aio() -> None:
llm = QianfanLLMEndpoint(streaming=True) # type: ignore[call-arg]
async for token in llm.astream("hi qianfan."):
assert isinstance(token, str)
def test_rate_limit() -> None:
llm = QianfanLLMEndpoint(model="ERNIE-Bot", init_kwargs={"query_per_second": 2}) # type: ignore[call-arg]
assert llm.client._client._rate_limiter._sync_limiter._query_per_second == 2
output = llm.generate(["write a joke"])
assert isinstance(output, LLMResult)
assert isinstance(output.generations, list)