seedream_image_task.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. import logging
  2. import time
  3. from app.config import settings
  4. from app.service.task.base_task import BaseTask
  5. from app.service.task.task_param import TaskParam
  6. from app.service.task.task_result import TaskResult
  7. from app.util.http_util import logged_client
  8. from app.util.retry_util import retry_on_failure
  9. logger = logging.getLogger(__name__)
  10. class SeedreamImageTask(BaseTask):
  11. def get_type(self) -> str:
  12. return "seedream_image"
  13. def _headers(self) -> dict[str, str]:
  14. return {
  15. "Content-Type": "application/json",
  16. "Authorization": f"Bearer {settings.seed_api_key}",
  17. }
  18. def _submit(self, url: str, inputs: dict) -> dict:
  19. with logged_client(timeout=60) as client:
  20. resp = client.post(url, headers=self._headers(), json=inputs)
  21. logger.info("[seedream_image] Submit response: status=%d", resp.status_code)
  22. resp.raise_for_status()
  23. data = resp.json()
  24. if "error" in data:
  25. err = data["error"]
  26. code = err.get("code", "")
  27. msg = err.get("message", "")
  28. raise RuntimeError(f"API error [{code}]: {msg}")
  29. return data
  30. def execute(self, param: TaskParam) -> TaskResult:
  31. url = settings.seedream_image_url
  32. logger.info("[seedream_image] Calling API: POST %s", url)
  33. t0 = time.time()
  34. try:
  35. data = retry_on_failure(
  36. lambda: self._submit(url, param.inputs),
  37. task_name="seedream_image",
  38. )
  39. images: list[str] = []
  40. for item in data.get("data", []):
  41. url_val = item.get("url")
  42. if url_val:
  43. images.append(url_val)
  44. logger.info("[seedream_image] Completed: model=%s, %d images, elapsed=%.2fs",
  45. data.get("model"), len(images), time.time() - t0)
  46. return TaskResult.success({"images": images})
  47. except Exception as e:
  48. elapsed = time.time() - t0
  49. logger.error("[seedream_image] Failed in %.2fs: %s", elapsed, e, exc_info=True)
  50. return TaskResult.failure(f"Failed to execute task: {e}")