158 lines
5.5 KiB
Python
158 lines
5.5 KiB
Python
import re
|
||
from typing import Any
|
||
|
||
IMAGE_STATUS_PENDING = "pending"
|
||
IMAGE_STATUS_PROCESSING = "processing"
|
||
IMAGE_STATUS_COMPLETED = "completed"
|
||
IMAGE_STATUS_FAILED = "failed"
|
||
|
||
VALID_IMAGE_STATUSES = {
|
||
IMAGE_STATUS_PENDING,
|
||
IMAGE_STATUS_PROCESSING,
|
||
IMAGE_STATUS_COMPLETED,
|
||
IMAGE_STATUS_FAILED,
|
||
}
|
||
|
||
_PLACEHOLDER_DESCRIPTION_RE = re.compile(
|
||
r"\{\{\{\{IMAGE:(.*?)\}\}\}\}|\{\{IMAGE:(.*?)\}\}"
|
||
)
|
||
|
||
|
||
def normalize_image_asset(asset: dict[str, Any] | None) -> dict[str, Any] | None:
|
||
"""归一化单条图片 dict。
|
||
|
||
- 兼容旧正文中的 {{{{IMAGE:…}}}} / {{IMAGE:…}} 占位符(需能解析出 description)。
|
||
- 新模型:插图不嵌入 markdown,可无占位符;已完成且带 url/storage_key 即可通过,
|
||
description 缺省时用「插图」;pending/processing 至少要有 description、占位符或 prompt。
|
||
"""
|
||
if not isinstance(asset, dict):
|
||
return None
|
||
|
||
ph_in = _as_non_empty_string(asset.get("placeholder"))
|
||
desc_in = _as_non_empty_string(asset.get("description"))
|
||
desc_from_ph = _extract_description_from_placeholder(ph_in) if ph_in else None
|
||
merged_description = desc_in or desc_from_ph
|
||
|
||
prompt_s = _as_non_empty_string(asset.get("prompt"))
|
||
error_s = _as_optional_string(asset.get("error"))
|
||
url_s = _as_optional_string(asset.get("url"))
|
||
storage_key_s = _as_optional_string(asset.get("storage_key"))
|
||
has_url_or_key = bool(url_s or storage_key_s)
|
||
|
||
status = _as_non_empty_string(asset.get("status")) or IMAGE_STATUS_PENDING
|
||
|
||
if ph_in and merged_description:
|
||
placeholder_out = ph_in
|
||
description_out = merged_description
|
||
elif status == IMAGE_STATUS_COMPLETED and has_url_or_key:
|
||
placeholder_out = ph_in or ""
|
||
description_out = merged_description or "插图"
|
||
elif status in (IMAGE_STATUS_PENDING, IMAGE_STATUS_PROCESSING):
|
||
if not (merged_description or ph_in or prompt_s):
|
||
return None
|
||
placeholder_out = ph_in or ""
|
||
description_out = merged_description or prompt_s or "插图"
|
||
elif status == IMAGE_STATUS_FAILED:
|
||
if not (merged_description or ph_in or error_s):
|
||
return None
|
||
placeholder_out = ph_in or ""
|
||
description_out = merged_description or "插图"
|
||
else:
|
||
return None
|
||
|
||
normalized = dict(asset)
|
||
normalized["index"] = _coerce_int(asset.get("index"), default=0)
|
||
normalized["placeholder"] = placeholder_out
|
||
normalized["description"] = description_out
|
||
if status not in VALID_IMAGE_STATUSES:
|
||
normalized["status"] = IMAGE_STATUS_FAILED
|
||
normalized["error"] = asset.get("error") or f"invalid image status: {status}"
|
||
return normalized
|
||
|
||
normalized["status"] = status
|
||
normalized["prompt"] = _as_optional_string(asset.get("prompt"))
|
||
normalized["url"] = _as_optional_string(asset.get("url"))
|
||
normalized["storage_key"] = _as_optional_string(asset.get("storage_key"))
|
||
normalized["provider"] = _as_optional_string(asset.get("provider"))
|
||
normalized["style"] = _as_optional_string(asset.get("style"))
|
||
normalized["size"] = _as_optional_string(asset.get("size"))
|
||
normalized["error"] = _as_optional_string(asset.get("error"))
|
||
normalized["created_at"] = _as_optional_string(asset.get("created_at"))
|
||
normalized["updated_at"] = _as_optional_string(asset.get("updated_at"))
|
||
normalized["retryable"] = _as_optional_bool(asset.get("retryable"))
|
||
|
||
if normalized["status"] == IMAGE_STATUS_COMPLETED and not (
|
||
normalized["url"] or normalized["storage_key"]
|
||
):
|
||
normalized["status"] = IMAGE_STATUS_FAILED
|
||
normalized["error"] = normalized["error"] or "missing image url"
|
||
|
||
if normalized["status"] != IMAGE_STATUS_FAILED:
|
||
normalized["retryable"] = None
|
||
|
||
return normalized
|
||
|
||
|
||
def normalize_image_assets(images: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
|
||
normalized_assets: list[dict[str, Any]] = []
|
||
for item in images or []:
|
||
normalized = normalize_image_asset(item)
|
||
if normalized:
|
||
normalized_assets.append(normalized)
|
||
return normalized_assets
|
||
|
||
|
||
def completed_image_assets(images: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
|
||
return [
|
||
asset
|
||
for asset in normalize_image_assets(images)
|
||
if asset.get("status") == IMAGE_STATUS_COMPLETED
|
||
and (asset.get("storage_key") or asset.get("url"))
|
||
]
|
||
|
||
|
||
def _as_non_empty_string(value: Any) -> str | None:
|
||
if isinstance(value, str):
|
||
stripped = value.strip()
|
||
return stripped or None
|
||
return None
|
||
|
||
|
||
def _as_optional_string(value: Any) -> str | None:
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, str):
|
||
return value
|
||
return str(value)
|
||
|
||
|
||
def _as_optional_bool(value: Any) -> bool | None:
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, bool):
|
||
return value
|
||
if isinstance(value, str):
|
||
lowered = value.strip().lower()
|
||
if lowered in {"true", "1", "yes"}:
|
||
return True
|
||
if lowered in {"false", "0", "no"}:
|
||
return False
|
||
return None
|
||
|
||
|
||
def _coerce_int(value: Any, default: int) -> int:
|
||
try:
|
||
return int(value)
|
||
except (TypeError, ValueError):
|
||
return default
|
||
|
||
|
||
def _extract_description_from_placeholder(placeholder: str | None) -> str | None:
|
||
if not placeholder:
|
||
return None
|
||
match = _PLACEHOLDER_DESCRIPTION_RE.fullmatch(placeholder.strip())
|
||
if not match:
|
||
return None
|
||
description = (match.group(1) or match.group(2) or "").strip()
|
||
return description or None
|