seedream_image_task.py 2.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  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. logger = logging.getLogger(__name__)
  9. class SeedreamImageTask(BaseTask):
  10. def get_type(self) -> str:
  11. return "seedream_image"
  12. def _headers(self) -> dict[str, str]:
  13. return {
  14. "Content-Type": "application/json",
  15. "Authorization": f"Bearer {settings.seed_api_key}",
  16. }
  17. def execute(self, param: TaskParam) -> TaskResult:
  18. url = settings.seedream_image_url
  19. logger.info("[seedream_image] Calling API: POST %s", url)
  20. t0 = time.time()
  21. try:
  22. with logged_client(timeout=60) as client:
  23. resp = client.post(url, headers=self._headers(), json=param.inputs)
  24. elapsed = time.time() - t0
  25. logger.info("[seedream_image] Response: status=%d, elapsed=%.2fs",
  26. resp.status_code, elapsed)
  27. data = resp.json()
  28. if "error" in data:
  29. err = data["error"]
  30. code = err.get("code", "")
  31. msg = err.get("message", "")
  32. logger.error("[seedream_image] API error: code=%s, message=%s", code, msg)
  33. return TaskResult.failure(f"API error [{code}]: {msg}")
  34. images: list[str] = []
  35. for item in data.get("data", []):
  36. url_val = item.get("url")
  37. if url_val:
  38. images.append(url_val)
  39. logger.info("[seedream_image] Completed: model=%s, %d images, elapsed=%.2fs",
  40. data.get("model"), len(images), time.time() - t0)
  41. return TaskResult.success({"images": images})
  42. except Exception as e:
  43. elapsed = time.time() - t0
  44. logger.error("[seedream_image] Failed in %.2fs: %s", elapsed, e, exc_info=True)
  45. return TaskResult.failure(f"Failed to execute task: {e}")