diff --git a/.env.example b/.env.example index 58f3793..98d29c7 100644 --- a/.env.example +++ b/.env.example @@ -15,6 +15,11 @@ OPENROUTER_BASE_URL=https://openrouter.ai/api/v1 OPENAI_API_KEY=your_openai_api_key # OPENAI_BASE_URL=https://api.openai.com/v1 +# Atlas Cloud (OpenAI-compatible) +ATLASCLOUD_API_KEY=your_atlascloud_api_key +ATLASCLOUD_API_BASE=https://api.atlascloud.ai/v1 +# ATLAS_CLOUD_API_KEY and ATLASCLOUD_BASE_URL are also supported aliases + # Anthropic ANTHROPIC_API_KEY=your_anthropic_api_key # ANTHROPIC_BASE_URL=https://api.anthropic.com diff --git a/mingli_bench/cli.py b/mingli_bench/cli.py index ef1e5c5..00026d5 100644 --- a/mingli_bench/cli.py +++ b/mingli_bench/cli.py @@ -117,7 +117,7 @@ def main(): parser.add_argument( "--platform", - choices=["openai", "openrouter", "anthropic", "google", "deepseek", "doubao"], + choices=["openai", "openrouter", "atlascloud", "anthropic", "google", "deepseek", "doubao"], help="Force routing platform (overrides auto-detection from model name prefix)", ) @@ -151,14 +151,14 @@ def main(): except Exception as e: logger.error(f"Failed to load dataset statistics: {e}") return 1 - print(f"\nDataset Statistics:") + print("\nDataset Statistics:") print(f" Name: {stats['benchmark_name']}") print(f" Version: {stats['data_version']}") print(f" Available Years: {', '.join(map(str, stats['available_years']))}") if args.year is not None: print(f" Selected Year: {args.year}") print(f" Total Questions: {stats['total_questions']}") - print(f"\n Categories:") + print("\n Categories:") for cat, count in stats['categories'].items(): print(f" - {cat}: {count}") return 0 @@ -216,7 +216,7 @@ def main(): if results['errors'] > 0: print(f"Errors: {results['errors']}") - print(f"\nCategory Breakdown:") + print("\nCategory Breakdown:") for cat, stats in results['category_stats'].items(): print(f" {cat:12s}: {stats['accuracy']:6.2%} ({stats['correct']}/{stats['total']})") diff --git a/mingli_bench/models/atlascloud_client.py b/mingli_bench/models/atlascloud_client.py new file mode 100644 index 0000000..3c75b1b --- /dev/null +++ b/mingli_bench/models/atlascloud_client.py @@ -0,0 +1,114 @@ +""" +Atlas Cloud model client implementation. +""" + +import os +from typing import Optional + +from openai import OpenAI + +from .base import ModelClient +from ..utils.logger import get_logger + +logger = get_logger(__name__) + + +class AtlasCloudClient(ModelClient): + """Client for Atlas Cloud models through the OpenAI-compatible API.""" + + API_KEY_ENV_VARS = ["ATLASCLOUD_API_KEY", "ATLAS_CLOUD_API_KEY"] + BASE_URL_ENV_VARS = [ + "ATLASCLOUD_API_BASE", + "ATLASCLOUD_BASE_URL", + "ATLAS_CLOUD_API_BASE", + "ATLAS_CLOUD_BASE_URL", + ] + DEFAULT_BASE_URL = "https://api.atlascloud.ai/v1" + + def __init__( + self, + model_name: str = "qwen/qwen3.5-flash", + api_key: Optional[str] = None, + base_url: Optional[str] = None, + **kwargs, + ): + """ + Initialize Atlas Cloud client. + + Args: + model_name: Atlas Cloud model id, e.g. "qwen/qwen3.5-flash" + api_key: Atlas Cloud API key + base_url: OpenAI-compatible base URL + **kwargs: Additional configuration + """ + super().__init__( + model_name=model_name, + api_key=api_key, + api_key_env_vars=self.API_KEY_ENV_VARS, + **kwargs, + ) + + timeout_seconds = int(os.getenv("TIMEOUT", "30")) + self.base_url = base_url or self._get_base_url() + self.client = OpenAI( + api_key=self.api_key, + base_url=self.base_url, + timeout=timeout_seconds, + ) + + logger.info( + "Initialized Atlas Cloud client with model: %s, base_url: %s, timeout: %ss", + model_name, + self.base_url, + timeout_seconds, + ) + + def _get_base_url(self) -> str: + """Resolve Atlas Cloud base URL from aliases or default.""" + for env_var in self.BASE_URL_ENV_VARS: + if value := os.getenv(env_var): + return value + return self.DEFAULT_BASE_URL + + def generate(self, prompt: str, **kwargs) -> str: + """ + Generate response using Atlas Cloud OpenAI-compatible API. + + Args: + prompt: Input prompt + **kwargs: Override generation parameters + + Returns: + Generated text + """ + try: + params = self.get_generation_params(**kwargs) + + response = self.client.chat.completions.create( + model=self.model_name, + messages=[ + {"role": "system", "content": self.SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + **params, + ) + + return response.choices[0].message.content.strip() + + except Exception as e: + self.handle_api_error("Atlas Cloud generation", e) + raise + + def validate_api_key(self) -> bool: + """ + Validate Atlas Cloud API key. + + Returns: + True if valid, False otherwise + """ + try: + self.client.models.list() + return True + except Exception as e: + logger.error(f"Invalid Atlas Cloud API key: {e}") + return False diff --git a/mingli_bench/models/factory.py b/mingli_bench/models/factory.py index ec55cf6..3f3c4ba 100644 --- a/mingli_bench/models/factory.py +++ b/mingli_bench/models/factory.py @@ -16,6 +16,7 @@ _PROVIDER_INSTALL_HINT = { 'openai': 'pip install openai', 'openrouter': 'pip install openai', + 'atlascloud': 'pip install openai', 'deepseek': 'pip install openai', 'anthropic': 'pip install anthropic', 'google': 'pip install google-generativeai', @@ -35,6 +36,7 @@ class ModelFactory: 'deepseek': ('.deepseek_client', 'DeepSeekClient'), 'doubao': ('.doubao_client', 'DoubaoClient'), 'openrouter': ('.openai_client', 'OpenAIClient'), # OpenAI-compatible API + 'atlascloud': ('.atlascloud_client', 'AtlasCloudClient'), # OpenAI-compatible API } @classmethod @@ -85,9 +87,11 @@ def get_provider(cls, model_name: str) -> Optional[str]: Returns: Provider name or None """ - # OpenRouter models (with prefix) + # OpenRouter / provider alias models (with prefix) if '/' in model_name: prefix = model_name.split('/')[0].lower() + if prefix in {'atlascloud', 'atlas-cloud', 'atlas'}: + return 'atlascloud' # Special case: bytedance prefix for Doubao native API if prefix == 'bytedance': return 'doubao' @@ -144,11 +148,15 @@ def create(cls, raise ValueError( f"Cannot determine provider for model '{model_name}'. " f"Supported patterns: gpt-*, o1-*, o3-*, o4-*, claude-*, gemini-*, deepseek-*, doubao-*, " - f"or use OpenRouter format: provider/model-name (e.g., openai/gpt-4, nvidia/llama-3)" + f"atlascloud/, or use OpenRouter format: provider/model-name " + f"(e.g., openai/gpt-4, nvidia/llama-3)" ) logger.info(f"Determined provider: {provider} for model: {model_name}") + if provider == 'atlascloud': + model_name = cls._normalize_atlascloud_model_name(model_name) + # Get provider configuration provider_config = config.get(provider, {}) @@ -183,6 +191,14 @@ def list_providers(cls) -> list: """Get list of available providers.""" return list(cls._registry.keys()) + @staticmethod + def _normalize_atlascloud_model_name(model_name: str) -> str: + """Strip routing aliases while preserving Atlas Cloud's canonical model id.""" + for prefix in ('atlascloud/', 'atlas-cloud/', 'atlas/'): + if model_name.lower().startswith(prefix): + return model_name[len(prefix):] + return model_name + @classmethod def list_supported_models(cls) -> Dict[str, list]: """Get supported models by provider.""" @@ -192,6 +208,7 @@ def list_supported_models(cls) -> Dict[str, list]: 'google': ['gemini-pro', 'gemini-1.5-pro', 'gemini-1.5-flash'], 'deepseek': ['deepseek-chat', 'deepseek-coder'], 'doubao': ['doubao-pro', 'doubao-lite'], + 'atlascloud': ['atlascloud/qwen/qwen3.5-flash', 'atlascloud/deepseek-ai/deepseek-v4-pro'], 'openrouter': [ 'openai/gpt-4', 'anthropic/claude-3-sonnet', 'google/gemini-2.0-flash', 'x-ai/grok-4', 'moonshotai/kimi-k2', 'deepseek/deepseek-r1' diff --git a/mingli_bench/utils/config.py b/mingli_bench/utils/config.py index 86a57eb..1b3b6cd 100644 --- a/mingli_bench/utils/config.py +++ b/mingli_bench/utils/config.py @@ -61,6 +61,20 @@ def load_config(env_file: Optional[str] = None) -> Dict[str, Any]: "max_tokens": default_max_tokens, }, + # Atlas Cloud OpenAI-compatible configuration + "atlascloud": { + "api_key": os.getenv("ATLASCLOUD_API_KEY") or os.getenv("ATLAS_CLOUD_API_KEY"), + "base_url": ( + os.getenv("ATLASCLOUD_API_BASE") + or os.getenv("ATLASCLOUD_BASE_URL") + or os.getenv("ATLAS_CLOUD_API_BASE") + or os.getenv("ATLAS_CLOUD_BASE_URL") + or "https://api.atlascloud.ai/v1" + ), + "temperature": default_temperature, + "max_tokens": default_max_tokens, + }, + # Native Anthropic configuration "anthropic": { "api_key": os.getenv("CLAUDE_API_KEY") or os.getenv("ANTHROPIC_API_KEY"), diff --git a/tests/test_atlascloud_provider.py b/tests/test_atlascloud_provider.py new file mode 100644 index 0000000..8810a09 --- /dev/null +++ b/tests/test_atlascloud_provider.py @@ -0,0 +1,91 @@ +import os +import unittest +from unittest.mock import patch + +from mingli_bench.models.atlascloud_client import AtlasCloudClient +from mingli_bench.models.factory import ModelFactory +from mingli_bench.utils.config import load_config + + +ATLAS_ENV_KEYS = [ + "ATLASCLOUD_API_KEY", + "ATLAS_CLOUD_API_KEY", + "ATLASCLOUD_API_BASE", + "ATLASCLOUD_BASE_URL", + "ATLAS_CLOUD_API_BASE", + "ATLAS_CLOUD_BASE_URL", +] + + +class TestAtlasCloudProvider(unittest.TestCase): + def setUp(self): + self.original_registry_entry = ModelFactory._registry["atlascloud"] + + def tearDown(self): + ModelFactory._registry["atlascloud"] = self.original_registry_entry + + def test_load_config_reads_atlascloud_aliases(self): + env = { + key: "" + for key in ATLAS_ENV_KEYS + } + env.update( + { + "ATLAS_CLOUD_API_KEY": "atlas-key", + "ATLASCLOUD_BASE_URL": "https://atlas.example/v1", + "MAX_TOKENS": "1024", + "TEMPERATURE": "0.2", + } + ) + + with patch.dict(os.environ, env, clear=False): + config = load_config(env_file="/tmp/mingli-bench-no-env-file") + + self.assertEqual(config["atlascloud"]["api_key"], "atlas-key") + self.assertEqual(config["atlascloud"]["base_url"], "https://atlas.example/v1") + self.assertEqual(config["atlascloud"]["max_tokens"], 1024) + self.assertEqual(config["atlascloud"]["temperature"], 0.2) + + def test_factory_detects_atlascloud_prefixes(self): + self.assertEqual(ModelFactory.get_provider("atlascloud/qwen/qwen3.5-flash"), "atlascloud") + self.assertEqual(ModelFactory.get_provider("atlas-cloud/deepseek-ai/deepseek-v4-pro"), "atlascloud") + self.assertEqual(ModelFactory.get_provider("atlas/qwen/qwen3.5-flash"), "atlascloud") + + def test_factory_strips_atlascloud_routing_prefix(self): + class FakeClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.model_name = kwargs["model_name"] + + ModelFactory._registry["atlascloud"] = FakeClient + + client = ModelFactory.create( + "atlascloud/qwen/qwen3.5-flash", + config={ + "atlascloud": { + "api_key": "atlas-key", + "base_url": "https://api.atlascloud.ai/v1", + } + }, + ) + + self.assertEqual(client.model_name, "qwen/qwen3.5-flash") + self.assertEqual(client.kwargs["api_key"], "atlas-key") + self.assertEqual(client.kwargs["base_url"], "https://api.atlascloud.ai/v1") + + def test_atlascloud_client_initializes_openai_compatible_client(self): + with patch.dict(os.environ, {"ATLASCLOUD_API_KEY": "atlas-key"}, clear=False): + with patch("mingli_bench.models.atlascloud_client.OpenAI") as openai_cls: + client = AtlasCloudClient(model_name="qwen/qwen3.5-flash") + + self.assertEqual(client.api_key, "atlas-key") + self.assertEqual(client.base_url, "https://api.atlascloud.ai/v1") + openai_cls.assert_called_once_with( + api_key="atlas-key", + base_url="https://api.atlascloud.ai/v1", + timeout=30, + ) + + +if __name__ == "__main__": + unittest.main()