From 84357523128b6c983543571cc026b013bf714b07 Mon Sep 17 00:00:00 2001 From: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> Date: Wed, 15 Jul 2026 09:58:01 +0800 Subject: [PATCH] Add Atlas Cloud provider --- astrbot/core/config/default.py | 13 ++++++ astrbot/core/provider/manager.py | 4 ++ .../provider/sources/atlascloud_source.py | 40 +++++++++++++++++++ tests/test_openai_source.py | 36 +++++++++++++++++ 4 files changed, 93 insertions(+) create mode 100644 astrbot/core/provider/sources/atlascloud_source.py diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index 6156601248..e2abc1678d 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -1352,6 +1352,19 @@ "proxy": "", "custom_headers": {}, }, + "Atlas Cloud": { + "id": "atlascloud", + "provider": "atlascloud", + "type": "atlascloud_chat_completion", + "provider_type": "chat_completion", + "enable": True, + "key": [], + "timeout": 120, + "api_base": "https://api.atlascloud.ai/v1", + "model": "qwen/qwen3.5-flash", + "proxy": "", + "custom_headers": {}, + }, "NVIDIA": { "id": "nvidia", "provider": "nvidia", diff --git a/astrbot/core/provider/manager.py b/astrbot/core/provider/manager.py index ae4001fcd6..3e73d99cab 100644 --- a/astrbot/core/provider/manager.py +++ b/astrbot/core/provider/manager.py @@ -387,6 +387,10 @@ def dynamic_import_provider(self, type: str) -> None: from .sources.openrouter_source import ( ProviderOpenRouter as ProviderOpenRouter, ) + case "atlascloud_chat_completion": + from .sources.atlascloud_source import ( + ProviderAtlasCloud as ProviderAtlasCloud, + ) case "anthropic_chat_completion": from .sources.anthropic_source import ( ProviderAnthropic as ProviderAnthropic, diff --git a/astrbot/core/provider/sources/atlascloud_source.py b/astrbot/core/provider/sources/atlascloud_source.py new file mode 100644 index 0000000000..96cf8bed96 --- /dev/null +++ b/astrbot/core/provider/sources/atlascloud_source.py @@ -0,0 +1,40 @@ +from ..register import register_provider_adapter +from .openai_source import ProviderOpenAIOfficial + +ATLASCLOUD_DEFAULT_API_BASE = "https://api.atlascloud.ai/v1" +ATLASCLOUD_DEFAULT_MODEL = "qwen/qwen3.5-flash" +ATLASCLOUD_MODELS = [ + "qwen/qwen3.5-flash", + "deepseek-ai/deepseek-v4-pro", + "deepseek-ai/deepseek-v4-flash", +] + + +@register_provider_adapter( + "atlascloud_chat_completion", + "Atlas Cloud Chat Completion Provider Adapter", +) +class ProviderAtlasCloud(ProviderOpenAIOfficial): + """Atlas Cloud provider using its OpenAI-compatible LLM endpoint.""" + + def __init__( + self, + provider_config: dict, + provider_settings: dict, + ) -> None: + if not provider_config.get("api_base"): + provider_config["api_base"] = ATLASCLOUD_DEFAULT_API_BASE + if not provider_config.get("model"): + provider_config["model"] = ATLASCLOUD_DEFAULT_MODEL + + super().__init__(provider_config, provider_settings) + + async def get_models(self) -> list[str]: + try: + models = await super().get_models() + if models: + return models + except Exception: + pass + + return ATLASCLOUD_MODELS.copy() diff --git a/tests/test_openai_source.py b/tests/test_openai_source.py index a45a232938..787ca22eb6 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -13,6 +13,7 @@ import astrbot.core.provider.sources.request_retry as request_retry from astrbot.core.exceptions import EmptyModelOutputError from astrbot.core.provider.entities import LLMResponse +from astrbot.core.provider.sources.atlascloud_source import ProviderAtlasCloud from astrbot.core.provider.sources.groq_source import ProviderGroq from astrbot.core.provider.sources.openai_source import ProviderOpenAIOfficial from astrbot.core.utils.media_utils import ResolvedMediaData, file_uri_to_path @@ -60,6 +61,22 @@ def _make_groq_provider(overrides: dict | None = None) -> ProviderGroq: ) +def _make_atlascloud_provider( + overrides: dict | None = None, +) -> ProviderAtlasCloud: + provider_config = { + "id": "test-atlascloud", + "type": "atlascloud_chat_completion", + "key": ["test-key"], + } + if overrides: + provider_config.update(overrides) + return ProviderAtlasCloud( + provider_config=provider_config, + provider_settings={}, + ) + + def test_create_http_client_uses_openai_httpx_module(monkeypatch): captured: dict[str, object] = {} @@ -120,6 +137,25 @@ def fake_import(name, globals=None, locals=None, fromlist=(), level=0): assert captured["httpx_module"] is openai_source_module.httpx +def test_atlascloud_provider_sets_default_openai_compatible_endpoint(): + provider = _make_atlascloud_provider() + + assert str(provider.client.base_url).rstrip("/") == "https://api.atlascloud.ai/v1" + assert provider.get_model() == "qwen/qwen3.5-flash" + + +def test_atlascloud_provider_keeps_custom_endpoint_and_model(): + provider = _make_atlascloud_provider( + { + "api_base": "https://proxy.example.com/v1", + "model": "deepseek-ai/deepseek-v4-pro", + }, + ) + + assert str(provider.client.base_url).rstrip("/") == "https://proxy.example.com/v1" + assert provider.get_model() == "deepseek-ai/deepseek-v4-pro" + + @pytest.mark.asyncio async def test_get_models_retries_transient_request_error(monkeypatch): monkeypatch.setattr(request_retry, "REQUEST_RETRY_WAIT_MIN_S", 0)