diff --git a/.env.example b/.env.example index 23edcdf91..b1fecb079 100644 --- a/.env.example +++ b/.env.example @@ -48,10 +48,13 @@ POE_API_KEY=your-poe-api-key-here #QWEN API KEY QWEN_API_KEY=your-qwen-key-here +# Atlas Cloud API Key +ATLASCLOUD_API_KEY=your-atlascloud-key-here + # Exa API Key (for Exa AI-powered web search MCP server) EXA_API_KEY=your-exa-api-key-here #Azure OpenAI Configuration AZURE_OPENAI_API_KEY=your-azure-openai-key-here AZURE_OPENAI_ENDPOINT=your-azure-openai-endpoint-here -AZURE_OPENAI_API_VERSION=your-azure-openai-api-version-here \ No newline at end of file +AZURE_OPENAI_API_VERSION=your-azure-openai-api-version-here diff --git a/docs/source/reference/supported_models.rst b/docs/source/reference/supported_models.rst index e3077dbef..022fc5766 100644 --- a/docs/source/reference/supported_models.rst +++ b/docs/source/reference/supported_models.rst @@ -156,7 +156,7 @@ The ``chatcompletion`` backend provides a generic way to connect to any OpenAI-c * - **Backend Type** - ``chatcompletion`` * - **Compatible Providers** - - Cerebras AI, Together AI, Fireworks AI, Groq, OpenRouter, POE, and any OpenAI-compatible API + - Cerebras AI, Together AI, Fireworks AI, Groq, OpenRouter, Atlas Cloud, POE, and any OpenAI-compatible API * - **Required Config** - ``base_url`` pointing to the provider's API endpoint * - **API Key** @@ -216,6 +216,9 @@ The ``chatcompletion`` backend provides a generic way to connect to any OpenAI-c * - **POE** - Platform-specific - Platform credentials + * - **Atlas Cloud** + - ``https://api.atlascloud.ai/v1`` + - ``ATLASCLOUD_API_KEY`` **Common Models:** @@ -223,6 +226,7 @@ The ``chatcompletion`` backend provides a generic way to connect to any OpenAI-c * **Together AI**: ``meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo``, ``mistralai/Mixtral-8x7B-Instruct-v0.1`` * **Fireworks AI**: ``accounts/fireworks/models/llama-v3p1-405b-instruct`` * **Groq**: ``llama-3.1-70b-versatile``, ``mixtral-8x7b-32768`` +* **Atlas Cloud**: use a model ID available to your Atlas Cloud account Tool Enablement Reference -------------------------- diff --git a/massgen/backend/chat_completions.py b/massgen/backend/chat_completions.py index 0208293a7..531ef0f2e 100644 --- a/massgen/backend/chat_completions.py +++ b/massgen/backend/chat_completions.py @@ -15,6 +15,7 @@ - ZAI: ZAI_API_KEY - POE: POE_API_KEY - Qwen: QWEN_API_KEY +- Atlas Cloud: ATLASCLOUD_API_KEY """ from __future__ import annotations @@ -32,6 +33,7 @@ from ..logger_config import log_backend_agent_message, log_stream_chunk, logger from ..stream_chunk import ChunkType from ..structured_logging import trace_llm_api_call +from ..utils.provider_urls import is_atlascloud_url # Local imports from ._constants import configure_openrouter_extra_body @@ -1145,6 +1147,8 @@ def get_provider_name(self) -> str: return "POE" elif "aliyuncs.com" in base_url: return "Qwen" + elif is_atlascloud_url(base_url): + return "Atlas Cloud" else: return "ChatCompletion" diff --git a/massgen/cli/backends.py b/massgen/cli/backends.py index d2bd793f5..e97b5d686 100644 --- a/massgen/cli/backends.py +++ b/massgen/cli/backends.py @@ -38,6 +38,7 @@ ) from ..logger_config import logger from ..utils import get_backend_type_from_model +from ..utils.provider_urls import is_atlascloud_url # --- cross-module references within the cli package --- from .config_loading import ( @@ -92,6 +93,7 @@ def create_backend(backend_type: str, **kwargs) -> Any: - Nvidia NIM (nvidia.com) -> NGC_API_KEY - POE (poe.com) -> POE_API_KEY - Qwen (dashscope.aliyuncs.com) -> QWEN_API_KEY + - Atlas Cloud (atlascloud.ai) -> ATLASCLOUD_API_KEY External agent frameworks are supported via the adapter registry. """ @@ -238,6 +240,15 @@ def create_backend(backend_type: str, **kwargs) -> Any: raise ConfigurationError( "Qwen API key not found. Set QWEN_API_KEY environment variable.\n" "You can add it to a .env file in:\n" " - Current directory: .env\n" " - Global config: ~/.massgen/.env", ) + elif is_atlascloud_url(base_url): + api_key = os.getenv("ATLASCLOUD_API_KEY") + if not api_key: + raise ConfigurationError( + "Atlas Cloud API key not found. Set ATLASCLOUD_API_KEY environment variable.\n" + "You can add it to a .env file in:\n" + " - Current directory: .env\n" + " - Global config: ~/.massgen/.env", + ) return ChatCompletionsBackend(api_key=api_key, **kwargs) diff --git a/massgen/configs/BACKEND_CONFIGURATION.md b/massgen/configs/BACKEND_CONFIGURATION.md index 0a31f0ad1..5e78a42e7 100644 --- a/massgen/configs/BACKEND_CONFIGURATION.md +++ b/massgen/configs/BACKEND_CONFIGURATION.md @@ -328,7 +328,7 @@ backend: Generic backend supporting multiple providers (v0.0.18+ with MCP). -Supports: Cerebras AI, Together AI, Fireworks AI, Groq, Nebius AI Studio, OpenRouter, Kimi/Moonshot, and any OpenAI-compatible API. +Supports: Cerebras AI, Together AI, Fireworks AI, Groq, Nebius AI Studio, OpenRouter, Kimi/Moonshot, Atlas Cloud, and any OpenAI-compatible API. ```yaml backend: @@ -365,6 +365,7 @@ backend: - **OpenRouter**: `https://openrouter.ai/api/v1` - **Nvidia NIM**: `https://integrate.api.nvidia.com/v1` - **Kimi/Moonshot**: `https://api.moonshot.cn/v1` +- **Atlas Cloud**: `https://api.atlascloud.ai/v1` --- @@ -464,4 +465,4 @@ NOTION_API_KEY=your-notion-key - [Join our Discord](https://discord.massgen.ai) - [Report Issues](https://github.com/Leezekun/MassGen/issues) -- [View Documentation](https://github.com/Leezekun/MassGen) \ No newline at end of file +- [View Documentation](https://github.com/Leezekun/MassGen) diff --git a/massgen/configs/basic/single/single_atlascloud.yaml b/massgen/configs/basic/single/single_atlascloud.yaml new file mode 100644 index 000000000..1c9562c00 --- /dev/null +++ b/massgen/configs/basic/single/single_atlascloud.yaml @@ -0,0 +1,16 @@ +# MassGen Atlas Cloud Configuration +# uv run massgen --automation --config massgen/configs/basic/single/single_atlascloud.yaml "Hello" +# Requires ATLASCLOUD_API_KEY in the environment. +# What happens: MassGen sends the prompt to Atlas Cloud through the +# OpenAI-compatible Chat Completions endpoint and renders the response in +# the textual terminal UI. +agents: + - id: "atlascloud_agent" + backend: + type: "chatcompletion" + model: "qwen/qwen3.8-max" + base_url: "https://api.atlascloud.ai/v1" + +ui: + display_type: "textual_terminal" + logging_enabled: true diff --git a/massgen/tests/test_chat_completions_refactor.py b/massgen/tests/test_chat_completions_refactor.py index 344e04e61..9befdc461 100644 --- a/massgen/tests/test_chat_completions_refactor.py +++ b/massgen/tests/test_chat_completions_refactor.py @@ -15,6 +15,9 @@ BACKGROUND_TOOL_STATUS_NAME, BACKGROUND_TOOL_WAIT_NAME, ) +from massgen.cli.backends import create_backend +from massgen.cli.config_loading import ConfigurationError +from massgen.utils.provider_urls import is_atlascloud_url def test_openai_backend_defaults(): @@ -47,6 +50,87 @@ def test_cerebras_backend(): assert backend.config["base_url"] == "https://api.cerebras.ai/v1" +@pytest.mark.parametrize( + ("base_url", "expected"), + [ + ("https://atlascloud.ai/v1", True), + ("https://api.atlascloud.ai/v1", True), + ("https://atlascloud.ai.evil.example/v1", False), + ("https://notatlascloud.ai/v1", False), + ("https://atlascloud.ai@evil.example/v1", False), + ("https://user@api.atlascloud.ai/v1", False), + ("https://evil.example/v1?next=https://api.atlascloud.ai/v1", False), + ], +) +def test_atlascloud_url_detection(base_url: str, expected: bool): + """Only exact Atlas Cloud hosts and valid subdomains should be detected.""" + assert is_atlascloud_url(base_url) is expected + + backend = ChatCompletionsBackend( + base_url=base_url, + api_key="test-key", + ) + expected_provider = "Atlas Cloud" if expected else "ChatCompletion" + assert backend.get_provider_name() == expected_provider + assert backend.config["base_url"] == base_url + + +@pytest.mark.parametrize( + "base_url", + ["https://atlascloud.ai/v1", "https://api.atlascloud.ai/v1"], +) +def test_atlascloud_api_key_is_detected( + monkeypatch: pytest.MonkeyPatch, + base_url: str, +): + """Atlas Cloud configs should use ATLASCLOUD_API_KEY automatically.""" + monkeypatch.setenv("ATLASCLOUD_API_KEY", "test-atlas-key") + + backend = create_backend( + "chatcompletion", + base_url=base_url, + model="qwen/qwen3.8-max", + ) + + assert backend.api_key == "test-atlas-key" + + +@pytest.mark.parametrize( + "base_url", + [ + "https://atlascloud.ai.evil.example/v1", + "https://atlascloud.ai@evil.example/v1", + "https://evil.example/v1?next=https://api.atlascloud.ai/v1", + ], +) +def test_atlascloud_api_key_is_not_used_for_untrusted_urls( + monkeypatch: pytest.MonkeyPatch, + base_url: str, +): + """Untrusted hosts must not receive the Atlas Cloud credential.""" + monkeypatch.setenv("ATLASCLOUD_API_KEY", "test-atlas-key") + + backend = create_backend( + "chatcompletion", + base_url=base_url, + model="qwen/qwen3.8-max", + ) + + assert backend.api_key is None + + +def test_atlascloud_api_key_error_is_actionable(monkeypatch: pytest.MonkeyPatch): + """Missing Atlas Cloud credentials should name the expected variable.""" + monkeypatch.delenv("ATLASCLOUD_API_KEY", raising=False) + + with pytest.raises(ConfigurationError, match="ATLASCLOUD_API_KEY"): + create_backend( + "chatcompletion", + base_url="https://api.atlascloud.ai/v1", + model="qwen/qwen3.8-max", + ) + + @pytest.mark.asyncio async def test_tool_conversion_via_api_params_handler(): """Response-style function tools are converted via ChatCompletionsAPIParamsHandler.""" diff --git a/massgen/utils/provider_urls.py b/massgen/utils/provider_urls.py new file mode 100644 index 000000000..9c881e75e --- /dev/null +++ b/massgen/utils/provider_urls.py @@ -0,0 +1,26 @@ +"""Provider URL validation helpers.""" + +from __future__ import annotations + +from urllib.parse import urlsplit + +_ATLASCLOUD_HOST = "atlascloud.ai" + + +def is_atlascloud_url(base_url: str | None) -> bool: + """Return whether a URL targets Atlas Cloud or one of its subdomains.""" + if not base_url: + return False + + try: + parsed = urlsplit(base_url) + hostname = parsed.hostname + has_userinfo = parsed.username is not None or parsed.password is not None + except (TypeError, ValueError): + return False + + if parsed.scheme.lower() not in {"http", "https"} or has_userinfo or not hostname: + return False + + hostname = hostname.rstrip(".").lower() + return hostname == _ATLASCLOUD_HOST or hostname.endswith(f".{_ATLASCLOUD_HOST}")