From 945f7d4bcf163e8651f21a6690a5184ec4acd6c9 Mon Sep 17 00:00:00 2001 From: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> Date: Wed, 22 Jul 2026 12:54:28 +0800 Subject: [PATCH] Add Atlas Cloud provider support --- backend/config.py | 2 +- backend/generators/atlascloud_image.py | 163 ++++++++++++++++++++++++ backend/generators/factory.py | 2 + backend/routes/config_routes.py | 63 ++++++++- backend/services/image.py | 8 ++ image_providers.yaml.example | 12 ++ tests/atlascloud_image_test.py | 151 ++++++++++++++++++++++ tests/errors_test.py | 54 ++++++++ tests/history_and_image_service_test.py | 32 +++++ text_providers.yaml.example | 7 + 10 files changed, 491 insertions(+), 3 deletions(-) create mode 100644 backend/generators/atlascloud_image.py create mode 100644 tests/atlascloud_image_test.py diff --git a/backend/config.py b/backend/config.py index 25cf6730..2d94ed67 100644 --- a/backend/config.py +++ b/backend/config.py @@ -133,7 +133,7 @@ def get_image_provider_config(cls, provider_name: str = None): ) provider_type = provider_config.get('type', provider_name) - if provider_type in ['openai', 'openai_compatible', 'image_api']: + if provider_type in ['openai', 'openai_compatible', 'image_api', 'atlascloud_image']: if not provider_config.get('base_url'): logger.error(f"服务商 [{provider_name}] 类型为 {provider_type}，但未配置 base_url") raise ValueError( diff --git a/backend/generators/atlascloud_image.py b/backend/generators/atlascloud_image.py new file mode 100644 index 00000000..fbaad817 --- /dev/null +++ b/backend/generators/atlascloud_image.py @@ -0,0 +1,163 @@ +"""Atlas Cloud Media API image generator.""" +import base64 +import logging +import time +from typing import Any, Dict, Optional + +import requests + +from .base import ImageGeneratorBase + +logger = logging.getLogger(__name__) + + +class AtlasCloudImageGenerator(ImageGeneratorBase): + """Generate images through the Atlas Cloud asynchronous Media API.""" + + DEFAULT_BASE_URL = "https://api.atlascloud.ai/api/v1" + DEFAULT_MODEL = "bytedance/seedream-v5.0-lite" + DEFAULT_SIZE = "1728*2304" + DEFAULT_OUTPUT_FORMAT = "png" + + def __init__(self, config: Dict[str, Any]): + super().__init__(config) + self.base_url = (config.get("base_url") or self.DEFAULT_BASE_URL).rstrip("/") + self.model = config.get("model") or self.DEFAULT_MODEL + self.size = config.get("size") or config.get("image_size") or self.DEFAULT_SIZE + self.output_format = config.get("output_format") or self.DEFAULT_OUTPUT_FORMAT + self.poll_interval_seconds = _non_negative_float( + config.get("poll_interval_seconds"), + default=3.0, + ) + self.max_poll_attempts = _positive_int( + config.get("max_poll_attempts"), + default=120, + ) + self.session = requests.Session() + + def validate_config(self) -> bool: + if not self.api_key: + raise ValueError( + "Atlas Cloud API Key 未配置。\n" + "解决方案：在系统设置页面编辑该服务商，填写 API Key" + ) + return True + + def generate_image( + self, + prompt: str, + model: Optional[str] = None, + **kwargs, + ) -> bytes: + """Submit an Atlas Cloud image task, poll for completion, then download the image.""" + self.validate_config() + + request_model = model or self.model + payload = { + "model": request_model, + "prompt": prompt, + "size": kwargs.get("size") or self.size, + "output_format": kwargs.get("output_format") or self.output_format, + "enable_base64_output": bool(kwargs.get("enable_base64_output", False)), + } + + logger.info(f"Atlas Cloud 生成图片: model={request_model}, size={payload['size']}") + prediction = self._submit_task(payload) + outputs = self._poll_outputs(prediction["id"]) + return self._read_output(outputs[0]) + + def _submit_task(self, payload: Dict[str, Any]) -> Dict[str, Any]: + response = self.session.post( + f"{self.base_url}/model/generateImage", + headers=self._headers(), + json=payload, + timeout=60, + ) + if response.status_code != 200: + raise Exception(f"Atlas Cloud 提交任务失败: HTTP {response.status_code}: {response.text[:500]}") + + data = _unwrap_data(response.json()) + request_id = data.get("id") or data.get("request_id") + if not request_id: + raise Exception(f"Atlas Cloud 提交响应缺少任务 ID: {str(data)[:500]}") + return {"id": request_id} + + def _poll_outputs(self, request_id: str) -> list[str]: + last_status = "" + for attempt in range(self.max_poll_attempts): + response = self._get_prediction_response(request_id) + if response.status_code != 200: + raise Exception(f"Atlas Cloud 查询任务失败: HTTP {response.status_code}: {response.text[:500]}") + + data = _unwrap_data(response.json()) + last_status = str(data.get("status") or "").lower() + outputs = data.get("outputs") or data.get("output") or [] + if isinstance(outputs, str): + outputs = [outputs] + + if last_status in {"completed", "succeeded", "success"} and outputs: + return outputs + if last_status in {"failed", "canceled", "cancelled"}: + raise Exception(f"Atlas Cloud 图片任务失败: {str(data)[:500]}") + + if attempt < self.max_poll_attempts - 1: + time.sleep(self.poll_interval_seconds) + + raise Exception(f"Atlas Cloud 图片任务超时，最后状态: {last_status or 'unknown'}") + + def _get_prediction_response(self, request_id: str): + response = self.session.get( + f"{self.base_url}/model/result/{request_id}", + headers=self._headers(), + timeout=60, + ) + if response.status_code != 404: + return response + return self.session.get( + f"{self.base_url}/model/prediction/{request_id}", + headers=self._headers(), + timeout=60, + ) + + def _read_output(self, output: str) -> bytes: + if output.startswith("data:image"): + return base64.b64decode(output.split(",", 1)[1]) + if _looks_like_base64(output): + return base64.b64decode(output) + + response = self.session.get(output, timeout=60) + if response.status_code != 200: + raise Exception(f"Atlas Cloud 图片下载失败: HTTP {response.status_code}") + return response.content + + def _headers(self) -> Dict[str, str]: + return { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + + +def _unwrap_data(payload: Dict[str, Any]) -> Dict[str, Any]: + data = payload.get("data") + return data if isinstance(data, dict) else payload + + +def _looks_like_base64(value: str) -> bool: + stripped = value.strip() + return bool(stripped) and not stripped.startswith(("http://", "https://")) + + +def _positive_int(value: Any, default: int) -> int: + try: + parsed = int(value) + return parsed if parsed > 0 else default + except (TypeError, ValueError): + return default + + +def _non_negative_float(value: Any, default: float) -> float: + try: + parsed = float(value) + return parsed if parsed >= 0 else default + except (TypeError, ValueError): + return default diff --git a/backend/generators/factory.py b/backend/generators/factory.py index 05b1d11c..839091ad 100644 --- a/backend/generators/factory.py +++ b/backend/generators/factory.py @@ -4,6 +4,7 @@ from .google_genai import GoogleGenAIGenerator from .openai_compatible import OpenAICompatibleGenerator from .image_api import ImageApiGenerator +from .atlascloud_image import AtlasCloudImageGenerator class ImageGeneratorFactory: @@ -15,6 +16,7 @@ class ImageGeneratorFactory: 'openai': OpenAICompatibleGenerator, 'openai_compatible': OpenAICompatibleGenerator, 'image_api': ImageApiGenerator, + 'atlascloud_image': AtlasCloudImageGenerator, } @classmethod diff --git a/backend/routes/config_routes.py b/backend/routes/config_routes.py index 2b071f7b..1996a475 100644 --- a/backend/routes/config_routes.py +++ b/backend/routes/config_routes.py @@ -132,7 +132,7 @@ def test_connection(): 测试服务商连接 请求体： - - type: 服务商类型（google_genai/google_gemini/openai_compatible/image_api） + - type: 服务商类型（google_genai/google_gemini/openai_compatible/image_api/atlascloud_image） - provider_name: 服务商名称（用于从配置读取 API Key） - api_key: API Key（可选，若不提供则从配置读取） - base_url: Base URL（可选） @@ -153,7 +153,7 @@ def test_connection(): return api_error_response( validation_error("缺少 type 参数", "请选择服务商类型后再测试连接。") ) - if provider_type not in ['google_genai', 'google_gemini', 'openai_compatible', 'image_api']: + if provider_type not in ['google_genai', 'google_gemini', 'openai_compatible', 'image_api', 'atlascloud_image']: return api_error_response( validation_error(f"不支持的类型: {provider_type}", "请选择正确的服务商类型后再测试连接。") ) @@ -324,6 +324,9 @@ def _test_provider_connection(provider_type: str, config: dict) -> dict: elif provider_type == 'image_api': return _test_image_api(config) + elif provider_type == 'atlascloud_image': + return _test_atlascloud_image(config) + else: raise ValueError(f"不支持的类型: {provider_type}") @@ -440,6 +443,62 @@ def _test_image_api(config: dict) -> dict: raise Exception(f"HTTP {response.status_code}: {response.text[:200]}") +def _test_atlascloud_image(config: dict) -> dict: + """测试 Atlas Cloud Media API 连接和模型目录。""" + import requests + + base_url = (config.get('base_url') or 'https://api.atlascloud.ai/api/v1').rstrip('/') + model = config.get('model') or 'bytedance/seedream-v5.0-lite' + + response = requests.get( + f"{base_url}/models", + headers={'Authorization': f"Bearer {config['api_key']}"}, + timeout=30, + ) + if response.status_code != 200: + raise Exception(f"HTTP {response.status_code}: {response.text[:200]}") + + try: + payload = response.json() + except Exception as exc: + raise Exception(f"Atlas Cloud 模型目录响应不是合法 JSON: {response.text[:500]}") from exc + + model_ids = _extract_model_ids(payload) + if model_ids and model not in model_ids: + return { + "success": True, + "warning": True, + "status": "warning", + "message": f"Atlas Cloud 连接成功，但模型目录中未找到 {model}。请确认模型名是否可用。" + } + + return { + "success": True, + "message": "Atlas Cloud 连接成功！模型目录可访问，图片任务提交将在生成时执行。" + } + + +def _extract_model_ids(payload) -> set[str]: + data = payload.get('data') if isinstance(payload, dict) else payload + if isinstance(data, dict): + for key in ['models', 'items', 'list']: + if isinstance(data.get(key), list): + data = data[key] + break + if not isinstance(data, list): + return set() + + model_ids = set() + for item in data: + if isinstance(item, dict): + model_id = item.get('model') or item.get('id') or item.get('name') + if isinstance(model_id, str): + model_ids.add(model_id) + elif isinstance(item, str): + model_ids.add(item) + return model_ids + + def _test_openai_chat_completion(config: dict, test_prompt: str) -> LlmSmokeResult: """用当前配置发送一次真实 OpenAI-compatible LLM 请求。""" import requests diff --git a/backend/services/image.py b/backend/services/image.py index e1c83181..ec91327d 100644 --- a/backend/services/image.py +++ b/backend/services/image.py @@ -214,6 +214,14 @@ def _generate_single_image( model=self.provider_config.get('model', 'nano-banana-2'), reference_images=reference_images if reference_images else None, ) + elif self.provider_config.get('type') == 'atlascloud_image': + logger.debug(" 使用 Atlas Cloud 图片生成器") + image_data = self.generator.generate_image( + prompt=prompt, + size=self.provider_config.get('size', '1728*2304'), + output_format=self.provider_config.get('output_format', 'png'), + model=self.provider_config.get('model', 'bytedance/seedream-v5.0-lite'), + ) else: logger.debug(f" 使用 OpenAI 兼容生成器") image_data = self.generator.generate_image( diff --git a/image_providers.yaml.example b/image_providers.yaml.example index 4d20a952..66acae7e 100644 --- a/image_providers.yaml.example +++ b/image_providers.yaml.example @@ -27,3 +27,15 @@ providers: base_url: https://your-api-endpoint.com model: dall-e-3 high_concurrency: false + + # Atlas Cloud Media API（异步图片生成） + atlascloud: + type: atlascloud_image + api_key: ak-xxxxxxxxxxxxxxxxxxxx + base_url: https://api.atlascloud.ai/api/v1 + model: bytedance/seedream-v5.0-lite + size: "1728*2304" + output_format: png + poll_interval_seconds: 3 + max_poll_attempts: 120 + high_concurrency: false diff --git a/tests/atlascloud_image_test.py b/tests/atlascloud_image_test.py new file mode 100644 index 00000000..dcd02cd0 --- /dev/null +++ b/tests/atlascloud_image_test.py @@ -0,0 +1,151 @@ +import yaml + +from backend.generators.atlascloud_image import AtlasCloudImageGenerator +from backend.generators.factory import ImageGeneratorFactory + + +class FakeResponse: + def __init__(self, status_code=200, payload=None, content=b"", text=""): + self.status_code = status_code + self._payload = payload or {} + self.content = content + self.text = text + + def json(self): + return self._payload + + +class FakeSession: + def __init__(self): + self.posts = [] + self.gets = [] + + def post(self, url, headers=None, json=None, timeout=None): + self.posts.append({ + "url": url, + "headers": headers, + "json": json, + "timeout": timeout, + }) + return FakeResponse(payload={ + "code": 200, + "data": { + "id": "prediction_123", + "status": "starting", + }, + }) + + def get(self, url, headers=None, timeout=None): + self.gets.append({ + "url": url, + "headers": headers, + "timeout": timeout, + }) + if url.endswith("/model/result/prediction_123"): + return FakeResponse(payload={ + "code": 200, + "data": { + "id": "prediction_123", + "status": "completed", + "outputs": ["https://cdn.example.com/generated.png"], + }, + }) + return FakeResponse(content=b"image-bytes") + + +class PredictionFallbackSession(FakeSession): + def get(self, url, headers=None, timeout=None): + self.gets.append({ + "url": url, + "headers": headers, + "timeout": timeout, + }) + if url.endswith("/model/result/prediction_123"): + return FakeResponse(status_code=404, text="not found") + if url.endswith("/model/prediction/prediction_123"): + return FakeResponse(payload={ + "data": { + "id": "prediction_123", + "status": "completed", + "outputs": ["data:image/png;base64,aW1hZ2UtYnl0ZXM="], + }, + }) + return FakeResponse(content=b"image-bytes") + + +def test_atlascloud_image_generator_submits_polls_and_downloads(): + session = FakeSession() + generator = AtlasCloudImageGenerator({ + "api_key": "ak-test", + "base_url": "https://api.atlascloud.ai/api/v1", + "model": "bytedance/seedream-v5.0-lite", + "size": "1728*2304", + "output_format": "png", + "poll_interval_seconds": 0, + "max_poll_attempts": 2, + }) + generator.session = session + + image = generator.generate_image("生成一张小红书封面") + + assert image == b"image-bytes" + assert session.posts[0]["url"] == "https://api.atlascloud.ai/api/v1/model/generateImage" + assert session.posts[0]["headers"]["Authorization"] == "Bearer ak-test" + assert session.posts[0]["json"] == { + "model": "bytedance/seedream-v5.0-lite", + "prompt": "生成一张小红书封面", + "size": "1728*2304", + "output_format": "png", + "enable_base64_output": False, + } + assert session.gets[0]["url"] == "https://api.atlascloud.ai/api/v1/model/result/prediction_123" + assert session.gets[1]["url"] == "https://cdn.example.com/generated.png" + + +def test_atlascloud_image_generator_falls_back_to_prediction_poll_path(): + session = PredictionFallbackSession() + generator = AtlasCloudImageGenerator({ + "api_key": "ak-test", + "base_url": "https://api.atlascloud.ai/api/v1", + "poll_interval_seconds": 0, + }) + generator.session = session + + image = generator.generate_image("生成一张小红书封面") + + assert image == b"image-bytes" + assert session.gets[0]["url"] == "https://api.atlascloud.ai/api/v1/model/result/prediction_123" + assert session.gets[1]["url"] == "https://api.atlascloud.ai/api/v1/model/prediction/prediction_123" + + +def test_atlascloud_image_generator_supports_data_url_output(): + generator = AtlasCloudImageGenerator({ + "api_key": "ak-test", + "poll_interval_seconds": 0, + }) + + assert generator._read_output("data:image/png;base64,aW1hZ2UtYnl0ZXM=") == b"image-bytes" + + +def test_factory_registers_atlascloud_image_generator(): + generator = ImageGeneratorFactory.create("atlascloud_image", {"api_key": "ak-test"}) + + assert isinstance(generator, AtlasCloudImageGenerator) + + +def test_atlascloud_provider_examples_are_valid_yaml(): + with open("image_providers.yaml.example", "r", encoding="utf-8") as f: + image_config = yaml.safe_load(f) + with open("text_providers.yaml.example", "r", encoding="utf-8") as f: + text_config = yaml.safe_load(f) + + image_provider = image_config["providers"]["atlascloud"] + text_provider = text_config["providers"]["atlascloud"] + + assert image_provider["type"] == "atlascloud_image" + assert image_provider["base_url"] == "https://api.atlascloud.ai/api/v1" + assert image_provider["model"] == "bytedance/seedream-v5.0-lite" + assert image_provider["size"] == "1728*2304" + assert text_provider["type"] == "openai_compatible" + assert text_provider["base_url"] == "https://api.atlascloud.ai/v1" + assert text_provider["model"] == "qwen/qwen3.5-flash" diff --git a/tests/errors_test.py b/tests/errors_test.py index 86d6d6ab..b0b437d1 100644 --- a/tests/errors_test.py +++ b/tests/errors_test.py @@ -303,6 +303,60 @@ def json(self): }) +def test_atlascloud_image_test_checks_model_catalog(monkeypatch): + import requests + from backend.routes import config_routes + + captured = {} + + class Response: + status_code = 200 + text = '{"data":[{"model":"bytedance/seedream-v5.0-lite"}]}' + + def json(self): + return {"data": [{"model": "bytedance/seedream-v5.0-lite"}]} + + def fake_get(url, headers=None, timeout=None): + captured["url"] = url + captured["headers"] = headers + captured["timeout"] = timeout + return Response() + + monkeypatch.setattr(requests, "get", fake_get) + + result = config_routes._test_provider_connection("atlascloud_image", { + "api_key": "ak-test", + "base_url": "https://api.atlascloud.ai/api/v1", + "model": "bytedance/seedream-v5.0-lite", + }) + + assert result["success"] is True + assert captured["url"] == "https://api.atlascloud.ai/api/v1/models" + assert captured["headers"]["Authorization"] == "Bearer ak-test" + assert captured["timeout"] == 30 + + +def test_config_test_accepts_atlascloud_image_type(client, monkeypatch): + from backend.routes import config_routes + + def fake_test(provider_type, config): + assert provider_type == "atlascloud_image" + assert config["api_key"] == "ak-test" + return {"success": True, "message": "ok"} + + monkeypatch.setattr(config_routes, "_test_provider_connection", fake_test) + + response = client.post("/api/config/test", json={ + "type": "atlascloud_image", + "api_key": "ak-test", + "base_url": "https://api.atlascloud.ai/api/v1", + "model": "bytedance/seedream-v5.0-lite", + }) + + assert response.status_code == 200 + assert response.get_json()["success"] is True + + def test_outline_missing_topic_returns_structured_error(client): response = client.post("/api/outline", json={"topic": ""}) data = response.get_json() diff --git a/tests/history_and_image_service_test.py b/tests/history_and_image_service_test.py index 58c5fed3..3b6bc822 100644 --- a/tests/history_and_image_service_test.py +++ b/tests/history_and_image_service_test.py @@ -89,7 +89,11 @@ def test_cached_generation_events_do_not_call_generator(tmp_path): class FakeGenerator: + def __init__(self): + self.calls = [] + def generate_image(self, **kwargs): + self.calls.append(kwargs) return b"image-bytes" @@ -124,6 +128,34 @@ def test_single_image_generation_writes_history_immediately(tmp_path): assert record["status"] == "completed" +def test_atlascloud_generation_uses_media_api_size(tmp_path): + image_service = ImageService.__new__(ImageService) + image_service.generator = FakeGenerator() + image_service.provider_config = { + "type": "atlascloud_image", + "model": "bytedance/seedream-v5.0-lite", + "size": "1728*2304", + "output_format": "png", + } + image_service.use_short_prompt = False + image_service.prompt_template = "{page_content}" + image_service.prompt_template_short = "" + image_service.current_task_dir = str(tmp_path / "task_1") + Path(image_service.current_task_dir).mkdir() + image_service.rate_limiter = ImageRateLimiter(max_concurrent=1, interval_seconds=0) + image_service.history_service = make_history_service(tmp_path) + + result = image_service._generate_single_image( + {"index": 0, "type": "cover", "content": "cover"}, + "task_1", + ) + + assert result == (0, True, "0.png", None) + assert image_service.generator.calls[0]["size"] == "1728*2304" + assert image_service.generator.calls[0]["output_format"] == "png" + assert image_service.generator.calls[0]["model"] == "bytedance/seedream-v5.0-lite" + + def test_retry_failed_images_creates_task_dir_and_merges_by_index(tmp_path): service = make_history_service(tmp_path) record_id = service.create_record("topic", { diff --git a/text_providers.yaml.example b/text_providers.yaml.example index 2b7fb31e..b448e47b 100644 --- a/text_providers.yaml.example +++ b/text_providers.yaml.example @@ -32,3 +32,10 @@ providers: api_key: sk-xxxxxxxxxxxxxxxxxxxx base_url: https://dashscope.aliyuncs.com/compatible-mode/v1 model: qwen-max + + # Atlas Cloud LLM API（OpenAI 兼容） + atlascloud: + type: openai_compatible + api_key: ak-xxxxxxxxxxxxxxxxxxxx + base_url: https://api.atlascloud.ai/v1 + model: qwen/qwen3.5-flash