dippoo Claude Opus 4.5 commited on
Commit
508b150
·
1 Parent(s): 7c6b44e

Fix WaveSpeed async job handling - add polling for outputs

Browse files

WaveSpeed returns async responses with urls.get for some models.
Added _poll_for_result method to poll until outputs are ready.
Also added _resolve_video_model helper.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

src/content_engine/services/cloud_providers/wavespeed_provider.py CHANGED
@@ -151,6 +151,62 @@ class WaveSpeedProvider(CloudProvider):
151
  return model_name
152
  return EDIT_MODEL_MAP["default"]
153
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
154
  @staticmethod
155
  def _ensure_min_image_size(image_bytes: bytes, min_pixels: int = 3686400) -> bytes:
156
  """Upscale image if total pixel count is below the minimum required by the API.
@@ -425,6 +481,11 @@ class WaveSpeedProvider(CloudProvider):
425
  image_url = out[0]
426
  elif isinstance(out, str):
427
  image_url = out
 
 
 
 
 
428
  elif isinstance(output, list) and output:
429
  image_url = output[0]
430
  elif isinstance(output, str):
 
151
  return model_name
152
  return EDIT_MODEL_MAP["default"]
153
 
154
+ def _resolve_video_model(self, model_name: str | None) -> str:
155
+ """Resolve a friendly name to a WaveSpeed video model API path."""
156
+ if model_name and model_name in VIDEO_MODEL_MAP:
157
+ return VIDEO_MODEL_MAP[model_name]
158
+ if model_name:
159
+ return model_name
160
+ return VIDEO_MODEL_MAP["default"]
161
+
162
+ async def _poll_for_result(self, poll_url: str, max_attempts: int = 60, interval: float = 2.0) -> str:
163
+ """Poll the WaveSpeed async job URL until outputs are ready.
164
+
165
+ Returns the first output URL when available.
166
+ """
167
+ import asyncio
168
+
169
+ for attempt in range(max_attempts):
170
+ try:
171
+ resp = await self._http_client.get(
172
+ poll_url,
173
+ headers={"Authorization": f"Bearer {self._api_key}"},
174
+ )
175
+ resp.raise_for_status()
176
+ result = resp.json()
177
+
178
+ data = result.get("data", result)
179
+ status = data.get("status", "")
180
+
181
+ if status == "failed":
182
+ error_msg = data.get("error", "Unknown error")
183
+ raise RuntimeError(f"WaveSpeed job failed: {error_msg}")
184
+
185
+ outputs = data.get("outputs", [])
186
+ if outputs:
187
+ logger.info("WaveSpeed job completed after %d polls", attempt + 1)
188
+ return outputs[0]
189
+
190
+ # Also check for 'output' field
191
+ if "output" in data:
192
+ out = data["output"]
193
+ if isinstance(out, list) and out:
194
+ return out[0]
195
+ elif isinstance(out, str):
196
+ return out
197
+
198
+ if status == "completed" and not outputs:
199
+ raise RuntimeError(f"WaveSpeed job completed but no outputs: {data}")
200
+
201
+ logger.debug("WaveSpeed job pending (attempt %d/%d)", attempt + 1, max_attempts)
202
+ await asyncio.sleep(interval)
203
+
204
+ except httpx.HTTPStatusError as e:
205
+ logger.warning("Poll request failed: %s", e)
206
+ await asyncio.sleep(interval)
207
+
208
+ raise RuntimeError(f"WaveSpeed job timed out after {max_attempts * interval}s")
209
+
210
  @staticmethod
211
  def _ensure_min_image_size(image_bytes: bytes, min_pixels: int = 3686400) -> bytes:
212
  """Upscale image if total pixel count is below the minimum required by the API.
 
481
  image_url = out[0]
482
  elif isinstance(out, str):
483
  image_url = out
484
+ # Async response: outputs empty but urls.get exists - need to poll
485
+ elif not outputs and "urls" in data and "get" in data.get("urls", {}):
486
+ poll_url = data["urls"]["get"]
487
+ logger.info("WaveSpeed returned async job, polling: %s", poll_url[:80])
488
+ image_url = await self._poll_for_result(poll_url)
489
  elif isinstance(output, list) and output:
490
  image_url = output[0]
491
  elif isinstance(output, str):