diff --git a/dataflow/serving/__init__.py b/dataflow/serving/__init__.py index df759981..e91abf95 100644 --- a/dataflow/serving/__init__.py +++ b/dataflow/serving/__init__.py @@ -1,4 +1,5 @@ from .api_llm_serving_request import APILLMServing_request +from .atlascloud_llm_serving import AtlasCloudLLMServing from .local_model_llm_serving import LocalModelLLMServing_vllm from .local_model_llm_serving import LocalModelLLMServing_sglang from .api_vlm_serving_openai import APIVLMServing_openai @@ -18,6 +19,7 @@ __all__ = [ "APIGoogleVertexAIServing", "APILLMServing_request", + "AtlasCloudLLMServing", "LocalModelLLMServing_vllm", "LocalModelLLMServing_sglang", "APIVLMServing_openai", diff --git a/dataflow/serving/atlascloud_llm_serving.py b/dataflow/serving/atlascloud_llm_serving.py new file mode 100644 index 00000000..e4c220e3 --- /dev/null +++ b/dataflow/serving/atlascloud_llm_serving.py @@ -0,0 +1,24 @@ +from .api_llm_serving_request import APILLMServing_request + + +class AtlasCloudLLMServing(APILLMServing_request): + """Atlas Cloud OpenAI-compatible LLM serving preset. + + Uses ``ATLASCLOUD_API_KEY`` by default and points to the Atlas Cloud + chat-completions endpoint while keeping the same request/response behavior + as ``APILLMServing_request``. + """ + + def __init__( + self, + api_url: str = "https://api.atlascloud.ai/v1/chat/completions", + key_name_of_api_key: str = "ATLASCLOUD_API_KEY", + model_name: str = "qwen/qwen3.5-flash", + **configs, + ): + super().__init__( + api_url=api_url, + key_name_of_api_key=key_name_of_api_key, + model_name=model_name, + **configs, + ) diff --git a/test/cpu_only/test_atlascloud_llm_serving.py b/test/cpu_only/test_atlascloud_llm_serving.py new file mode 100644 index 00000000..62c17fcf --- /dev/null +++ b/test/cpu_only/test_atlascloud_llm_serving.py @@ -0,0 +1,30 @@ +import pytest + +from dataflow.serving import AtlasCloudLLMServing + + +def test_atlascloud_defaults(monkeypatch): + monkeypatch.setenv("ATLASCLOUD_API_KEY", "dummy-key") + + serving = AtlasCloudLLMServing(max_workers=1, max_retries=1) + + assert serving.api_url == "https://api.atlascloud.ai/v1/chat/completions" + assert serving.model_name == "qwen/qwen3.5-flash" + assert serving.api_key == "dummy-key" + serving.cleanup() + + +def test_atlascloud_uses_custom_model(monkeypatch): + monkeypatch.setenv("ATLASCLOUD_API_KEY", "dummy-key") + + serving = AtlasCloudLLMServing(model_name="deepseek-ai/deepseek-v4-pro") + + assert serving.model_name == "deepseek-ai/deepseek-v4-pro" + serving.cleanup() + + +def test_atlascloud_requires_api_key(monkeypatch): + monkeypatch.delenv("ATLASCLOUD_API_KEY", raising=False) + + with pytest.raises(ValueError, match="ATLASCLOUD_API_KEY"): + AtlasCloudLLMServing()