diff --git a/kie_api/jobs.py b/kie_api/jobs.py index e69de29..7624e19 100644 --- a/kie_api/jobs.py +++ b/kie_api/jobs.py @@ -0,0 +1,166 @@ +"""Shared job lifecycle helpers for KIE async tasks. + +This module centralizes common KIE job lifecycle operations that are model-agnostic: + +- Fetch `recordInfo` for a task id. +- Decide whether a task failure is likely transient. +- Poll a task until completion, failure, or timeout. + +Behavior (logging text, timing, error types) is intentionally kept identical to the +original model-specific implementations. +""" + +import json +import time +from typing import Any + +from .http import TransientKieError, requests + + +RECORD_INFO_URL = "https://api.kie.ai/api/v1/jobs/recordInfo" + + +def _log(enabled: bool, msg: str) -> None: + if enabled: + print(f"[KIE] {msg}") + + +def _fetch_task_record(api_key: str, task_id: str) -> tuple[dict[str, Any], str, Any]: + """Fetch the current record for a task. + + Returns: + A tuple of (data_dict, raw_response_text, message_field). + Raises: + RuntimeError: If the request fails, returns non-JSON, or returns an error code. + TransientKieError: If the API responds with retryable errors (429 or >=500). + """ + try: + response = requests.get( + RECORD_INFO_URL, + headers={"Authorization": f"Bearer {api_key}"}, + params={"taskId": task_id}, + timeout=30, + ) + except requests.RequestException as exc: + raise RuntimeError(f"Failed to call recordInfo endpoint: {exc}") from exc + + if response.status_code == 429 or response.status_code >= 500: + raise TransientKieError( + f"recordInfo returned HTTP {response.status_code}: {response.text}", status_code=response.status_code + ) + + raw_text = response.text + try: + payload_json: Any = response.json() + except json.JSONDecodeError as exc: + raise RuntimeError("recordInfo endpoint did not return valid JSON.") from exc + + if payload_json.get("code") != 200: + message = payload_json.get("message") or payload_json.get("msg") + raise RuntimeError(f"recordInfo endpoint returned error code {payload_json.get('code')}: {message}") + + data = payload_json.get("data") + if data is None: + raise RuntimeError("recordInfo endpoint returned no data field.") + + return data, raw_text, payload_json.get("message") or payload_json.get("msg") + + +def _should_retry_fail(fail_code: Any, fail_msg: Any, message: Any) -> bool: + """Determine whether a failed task should be retried. + + Returns: + True if the failure looks transient and worth retrying, otherwise False. + Raises: + None. + """ + try: + code_int = int(fail_code) + except (TypeError, ValueError): + code_int = None + + if code_int is not None and 500 <= code_int <= 599: + return True + + combined_text = " ".join( + str(part).lower() + for part in (fail_msg, message) + if isinstance(part, str) and part + ) + + if "internal error" in combined_text or "try again later" in combined_text: + return True + + return False + + +def _poll_task_until_complete( + api_key: str, + task_id: str, + poll_interval_s: float, + timeout_s: int, + log: bool, + start_time: float, +) -> dict[str, Any]: + """Poll recordInfo until the task completes, fails, or times out. + + Returns: + The task record data dict returned by recordInfo. + Raises: + RuntimeError: If the task times out or returns a non-retryable failure. + TransientKieError: If the task fails with a retryable condition. + """ + # Ensure we never poll faster than once per second to reduce server load. + interval = poll_interval_s if poll_interval_s > 0 else 1.0 + last_state = None + last_log_time = start_time + + while True: + now = time.time() + elapsed = now - start_time + if elapsed > timeout_s: + last_state_text = last_state if last_state is not None else "unknown" + raise RuntimeError( + f"Task {task_id} timed out after {timeout_s}s (last state={last_state_text}, elapsed={elapsed:.1f}s). " + "Try increasing timeout or retry." + ) + + data, raw_json, message_field = _fetch_task_record(api_key, task_id) + state = data.get("state") + # Log only on state change or every 30s to give progress without noisy output. + should_log = log and (state != last_state or (now - last_log_time) >= 30.0) + if should_log: + _log( + log, + f"Task {task_id} state: {state or 'unknown'} " + f"(elapsed={elapsed:.1f}s)" + ) + last_log_time = now + + last_state = state + + if state == "success": + if log: + _log(log, f"Task {task_id} completed (elapsed={elapsed:.1f}s)") + return data + if state == "fail": + fail_code = data.get("failCode") + fail_msg = data.get("failMsg") or data.get("msg") + parts = [f"Task {task_id} failed"] + if fail_code is not None: + parts.append(f"failCode={fail_code}") + if fail_msg: + parts.append(f"failMsg={fail_msg}") + if message_field: + parts.append(f"message={message_field}") + error_message = "; ".join(parts) + + if _should_retry_fail(fail_code, fail_msg, message_field): + raise TransientKieError(error_message) + + raise RuntimeError(error_message) + + if should_log: + _log(log, f"Polling again in {interval} seconds...") + time.sleep(interval) + diff --git a/kie_api/nanobanana.py b/kie_api/nanobanana.py index 6c93c10..6ed8f84 100644 --- a/kie_api/nanobanana.py +++ b/kie_api/nanobanana.py @@ -21,6 +21,7 @@ from PIL import Image from .auth import _load_api_key from .credits import _log_remaining_credits from .http import TransientKieError, requests +from .jobs import _fetch_task_record, _poll_task_until_complete, _should_retry_fail from .upload import _image_tensor_to_png_bytes, _truncate_url, _upload_image from .images import _image_bytes_to_tensor @@ -120,144 +121,6 @@ def _create_nano_banana_task(api_key: str, payload: dict[str, Any]) -> tuple[str return task_id, response.text -def _fetch_task_record(api_key: str, task_id: str) -> tuple[dict[str, Any], str, Any]: - """Fetch the current record for a task. - - Returns: - A tuple of (data_dict, raw_response_text, message_field). - Raises: - RuntimeError: If the request fails, returns non-JSON, or returns an error code. - TransientKieError: If the API responds with retryable errors (429 or >=500). - """ - try: - response = requests.get( - RECORD_INFO_URL, - headers={"Authorization": f"Bearer {api_key}"}, - params={"taskId": task_id}, - timeout=30, - ) - except requests.RequestException as exc: - raise RuntimeError(f"Failed to call recordInfo endpoint: {exc}") from exc - - if response.status_code == 429 or response.status_code >= 500: - raise TransientKieError( - f"recordInfo returned HTTP {response.status_code}: {response.text}", status_code=response.status_code - ) - - raw_text = response.text - try: - payload_json: Any = response.json() - except json.JSONDecodeError as exc: - raise RuntimeError("recordInfo endpoint did not return valid JSON.") from exc - - if payload_json.get("code") != 200: - message = payload_json.get("message") or payload_json.get("msg") - raise RuntimeError(f"recordInfo endpoint returned error code {payload_json.get('code')}: {message}") - - data = payload_json.get("data") - if data is None: - raise RuntimeError("recordInfo endpoint returned no data field.") - - return data, raw_text, payload_json.get("message") or payload_json.get("msg") - - -def _should_retry_fail(fail_code: Any, fail_msg: Any, message: Any) -> bool: - """Determine whether a failed task should be retried. - - Returns: - True if the failure looks transient and worth retrying, otherwise False. - Raises: - None. - """ - try: - code_int = int(fail_code) - except (TypeError, ValueError): - code_int = None - - if code_int is not None and 500 <= code_int <= 599: - return True - - combined_text = " ".join( - str(part).lower() - for part in (fail_msg, message) - if isinstance(part, str) and part - ) - - if "internal error" in combined_text or "try again later" in combined_text: - return True - - return False - - -def _poll_task_until_complete( - api_key: str, - task_id: str, - poll_interval_s: float, - timeout_s: int, - log: bool, - start_time: float, -) -> dict[str, Any]: - """Poll recordInfo until the task completes, fails, or times out. - - Returns: - The task record data dict returned by recordInfo. - Raises: - RuntimeError: If the task times out or returns a non-retryable failure. - TransientKieError: If the task fails with a retryable condition. - """ - # Ensure we never poll faster than once per second to reduce server load. - interval = poll_interval_s if poll_interval_s > 0 else 1.0 - last_state = None - last_log_time = start_time - - while True: - now = time.time() - elapsed = now - start_time - if elapsed > timeout_s: - last_state_text = last_state if last_state is not None else "unknown" - raise RuntimeError( - f"Task {task_id} timed out after {timeout_s}s (last state={last_state_text}, elapsed={elapsed:.1f}s). " - "Try increasing timeout or retry." - ) - - data, raw_json, message_field = _fetch_task_record(api_key, task_id) - state = data.get("state") - # Log only on state change or every 30s to give progress without noisy output. - should_log = log and (state != last_state or (now - last_log_time) >= 30.0) - if should_log: - _log( - log, - f"Task {task_id} state: {state or 'unknown'} " - f"(elapsed={elapsed:.1f}s)" - ) - last_log_time = now - - last_state = state - - if state == "success": - if log: - _log(log, f"Task {task_id} completed (elapsed={elapsed:.1f}s)") - return data - if state == "fail": - fail_code = data.get("failCode") - fail_msg = data.get("failMsg") or data.get("msg") - parts = [f"Task {task_id} failed"] - if fail_code is not None: - parts.append(f"failCode={fail_code}") - if fail_msg: - parts.append(f"failMsg={fail_msg}") - if message_field: - parts.append(f"message={message_field}") - error_message = "; ".join(parts) - - if _should_retry_fail(fail_code, fail_msg, message_field): - raise TransientKieError(error_message) - - raise RuntimeError(error_message) - - if should_log: - _log(log, f"Polling again in {interval} seconds...") - time.sleep(interval) def _create_nanobanana_task(api_key: str, payload: dict[str, Any]) -> tuple[str, str]: diff --git a/kie_api/seedream45_t2i.py b/kie_api/seedream45_t2i.py index 8716de1..55d52ce 100644 --- a/kie_api/seedream45_t2i.py +++ b/kie_api/seedream45_t2i.py @@ -10,11 +10,11 @@ from .auth import _load_api_key from .credits import _log_remaining_credits from .http import TransientKieError, requests from .images import _image_bytes_to_tensor +from .jobs import _poll_task_until_complete from .nanobanana import ( _download_image, _extract_result_urls, _log, - _poll_task_until_complete, )