| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061 |
- import logging
- import time
- from app.config import settings
- from app.service.task.base_task import BaseTask
- from app.service.task.task_param import TaskParam
- from app.service.task.task_result import TaskResult
- from app.util.http_util import logged_client
- from app.util.retry_util import retry_on_failure
- logger = logging.getLogger(__name__)
- class SeedreamImageTask(BaseTask):
- def get_type(self) -> str:
- return "seedream_image"
- def _headers(self) -> dict[str, str]:
- return {
- "Content-Type": "application/json",
- "Authorization": f"Bearer {settings.seed_api_key}",
- }
- def _submit(self, url: str, inputs: dict) -> dict:
- with logged_client(timeout=60) as client:
- resp = client.post(url, headers=self._headers(), json=inputs)
- logger.info("[seedream_image] Submit response: status=%d", resp.status_code)
- if not resp.is_success:
- raise RuntimeError(f"HTTP {resp.status_code}: {resp.text}")
- data = resp.json()
- if "error" in data:
- err = data["error"]
- code = err.get("code", "")
- msg = err.get("message", "")
- raise RuntimeError(f"API error [{code}]: {msg}")
- return data
- def execute(self, param: TaskParam) -> TaskResult:
- url = settings.seedream_image_url
- logger.info("[seedream_image] Calling API: POST %s", url)
- t0 = time.time()
- try:
- data = retry_on_failure(
- lambda: self._submit(url, param.inputs),
- task_name="seedream_image",
- )
- images: list[str] = []
- for item in data.get("data", []):
- url_val = item.get("url")
- if url_val:
- images.append(url_val)
- logger.info("[seedream_image] Completed: model=%s, %d images, elapsed=%.2fs",
- data.get("model"), len(images), time.time() - t0)
- return TaskResult.success({"images": images})
- except Exception as e:
- elapsed = time.time() - t0
- logger.error("[seedream_image] Failed in %.2fs: %s", elapsed, e, exc_info=True)
- return TaskResult.failure(f"Failed to execute task: {e}")
|
