121 lines
4.3 KiB
Python
121 lines
4.3 KiB
Python
"""Tests for Chinese LLM Providers (Wenxin, Doubao, Yuanbao)"""
|
|
|
|
import pytest
|
|
|
|
from agentkit.llm.providers.wenxin import WenxinProvider, WENXIN_MODEL_MAP
|
|
from agentkit.llm.providers.doubao import DoubaoProvider, DOUBAO_MODEL_MAP
|
|
from agentkit.llm.providers.yuanbao import YuanbaoProvider, YUANBAO_MODEL_MAP
|
|
from agentkit.llm.protocol import LLMRequest
|
|
|
|
|
|
class TestWenxinProvider:
|
|
"""WenxinProvider unit tests"""
|
|
|
|
def test_init_with_api_key(self):
|
|
provider = WenxinProvider(api_key="test_key")
|
|
assert provider._api_key == "test_key"
|
|
assert provider._default_model == "ernie-4.5-turbo-128k"
|
|
|
|
def test_init_with_ak_sk(self):
|
|
provider = WenxinProvider(
|
|
api_key="",
|
|
access_key="test_ak",
|
|
secret_key="test_sk",
|
|
)
|
|
assert provider._access_key == "test_ak"
|
|
assert provider._secret_key == "test_sk"
|
|
|
|
def test_model_mapping(self):
|
|
assert "ernie-4.5-turbo-128k" in WENXIN_MODEL_MAP
|
|
assert "ernie-5.0" in WENXIN_MODEL_MAP
|
|
assert "ernie-x1.1" in WENXIN_MODEL_MAP
|
|
|
|
def test_default_base_url(self):
|
|
from agentkit.llm.providers.wenxin import WENXIN_DEFAULT_BASE_URL
|
|
assert "qianfan.baidubce.com" in WENXIN_DEFAULT_BASE_URL
|
|
|
|
def test_custom_base_url(self):
|
|
provider = WenxinProvider(api_key="test", base_url="https://custom.api.com/v2")
|
|
assert "custom.api.com" in provider._base_url
|
|
|
|
|
|
class TestDoubaoProvider:
|
|
"""DoubaoProvider unit tests"""
|
|
|
|
def test_init(self):
|
|
provider = DoubaoProvider(api_key="test_key")
|
|
assert provider._api_key == "test_key"
|
|
assert provider._default_model == "doubao-pro-32k"
|
|
|
|
def test_model_mapping(self):
|
|
assert "doubao-pro-32k" in DOUBAO_MODEL_MAP
|
|
assert "doubao-lite-32k" in DOUBAO_MODEL_MAP
|
|
assert "doubao-vision" in DOUBAO_MODEL_MAP
|
|
|
|
def test_default_base_url(self):
|
|
from agentkit.llm.providers.doubao import DOUBAO_DEFAULT_BASE_URL
|
|
assert "ark.cn-beijing.volces.com" in DOUBAO_DEFAULT_BASE_URL
|
|
|
|
def test_custom_model(self):
|
|
provider = DoubaoProvider(
|
|
api_key="test",
|
|
default_model="doubao-lite-32k",
|
|
)
|
|
assert provider._default_model == "doubao-lite-32k"
|
|
|
|
|
|
class TestYuanbaoProvider:
|
|
"""YuanbaoProvider unit tests"""
|
|
|
|
def test_init(self):
|
|
provider = YuanbaoProvider(api_key="test_key")
|
|
assert provider._api_key == "test_key"
|
|
assert provider._default_model == "hunyuan-turbos-latest"
|
|
|
|
def test_init_with_enhancement(self):
|
|
provider = YuanbaoProvider(api_key="test", enable_enhancement=True)
|
|
assert provider._enable_enhancement is True
|
|
|
|
def test_model_mapping(self):
|
|
assert "hunyuan-turbos-latest" in YUANBAO_MODEL_MAP
|
|
assert "hunyuan-2.0" in YUANBAO_MODEL_MAP
|
|
assert "hunyuan-t1" in YUANBAO_MODEL_MAP
|
|
|
|
def test_default_base_url(self):
|
|
from agentkit.llm.providers.yuanbao import YUANBAO_DEFAULT_BASE_URL
|
|
assert "hunyuan.cloud.tencent.com" in YUANBAO_DEFAULT_BASE_URL
|
|
|
|
def test_enhancement_disabled_by_default(self):
|
|
provider = YuanbaoProvider(api_key="test")
|
|
assert provider._enable_enhancement is False
|
|
|
|
|
|
class TestProviderImports:
|
|
"""Test that all providers are importable from the package"""
|
|
|
|
def test_import_all_providers(self):
|
|
from agentkit.llm.providers import (
|
|
AnthropicProvider,
|
|
DoubaoProvider,
|
|
GeminiProvider,
|
|
OpenAICompatibleProvider,
|
|
WenxinProvider,
|
|
YuanbaoProvider,
|
|
)
|
|
assert AnthropicProvider is not None
|
|
assert DoubaoProvider is not None
|
|
assert GeminiProvider is not None
|
|
assert OpenAICompatibleProvider is not None
|
|
assert WenxinProvider is not None
|
|
assert YuanbaoProvider is not None
|
|
|
|
def test_inheritance(self):
|
|
"""All providers should inherit from OpenAICompatibleProvider or LLMProvider"""
|
|
from agentkit.llm.providers.openai import OpenAICompatibleProvider
|
|
from agentkit.llm.protocol import LLMProvider
|
|
|
|
assert issubclass(WenxinProvider, OpenAICompatibleProvider)
|
|
assert issubclass(DoubaoProvider, OpenAICompatibleProvider)
|
|
assert issubclass(YuanbaoProvider, OpenAICompatibleProvider)
|
|
assert issubclass(OpenAICompatibleProvider, LLMProvider)
|