from __future__ import annotations from dataclasses import dataclass from datetime import datetime, timedelta, timezone from app.config import settings from app.enums.generation_task import GenerationType from app.models.chat_generation_task import ChatGenerationTask from app.services.redis_registry_service import ensure_aware_utc @dataclass(slots=True) class PollScheduleDecision: delay_seconds: int next_poll_at: datetime poll_interval_seconds: int direct_countdown: bool final_poll: bool reason: str def utc_now() -> datetime: return datetime.now(timezone.utc) def is_video_generation_task(task: ChatGenerationTask) -> bool: return str(getattr(task, "gen_type", "") or "").lower() == GenerationType.VIDEO.value def video_final_deadline_from(now: datetime | None = None) -> datetime: current_time = now or utc_now() hours = max(1, int(settings.CHATAPI_ASYNC_VIDEO_FINAL_DEADLINE_HOURS or 24)) return current_time + timedelta(hours=hours) def ensure_video_poll_fields(task: ChatGenerationTask, *, now: datetime | None = None) -> None: """补齐视频轮询调度字段,兼容历史任务。""" if not is_video_generation_task(task): return current_time = now or utc_now() if ensure_aware_utc(getattr(task, "poll_started_at", None)) is None: task.poll_started_at = current_time if ensure_aware_utc(getattr(task, "deadline_at", None)) is None: task.deadline_at = video_final_deadline_from(current_time) if getattr(task, "poll_interval_seconds", None) is None: task.poll_interval_seconds = 0 def is_final_poll_due(task: ChatGenerationTask, *, now: datetime | None = None) -> bool: deadline_at = ensure_aware_utc(getattr(task, "deadline_at", None)) return bool(deadline_at and deadline_at <= (now or utc_now())) def is_poll_not_due(task: ChatGenerationTask, *, now: datetime | None = None, tolerance_seconds: int = 1) -> bool: """判断当前 poll 任务是否早于 next_poll_at。只对视频降频轮询生效。""" if not is_video_generation_task(task): return False if is_final_poll_due(task, now=now): return False next_poll_at = ensure_aware_utc(getattr(task, "next_poll_at", None)) if next_poll_at is None: return False return next_poll_at > (now or utc_now()) + timedelta(seconds=max(0, int(tolerance_seconds))) def _clamp_positive_seconds(value: int | float | None, default: int) -> int: try: parsed = int(value if value is not None else default) except (TypeError, ValueError): parsed = default return max(1, parsed) def build_video_pending_poll_schedule( task: ChatGenerationTask, *, now: datetime | None = None, ) -> PollScheduleDecision: """计算视频任务 pending/running 后的下一次轮询时间。""" current_time = now or utc_now() ensure_video_poll_fields(task, now=current_time) deadline_at = ensure_aware_utc(task.deadline_at) if deadline_at and deadline_at <= current_time: return PollScheduleDecision( delay_seconds=0, next_poll_at=current_time, poll_interval_seconds=int(task.poll_interval_seconds or 0), direct_countdown=True, final_poll=True, reason="video_final_poll_due", ) poll_started_at = ensure_aware_utc(task.poll_started_at) or current_time elapsed_seconds = max(0, int((current_time - poll_started_at).total_seconds())) high_freq_seconds = max(0, int(settings.CHATAPI_ASYNC_VIDEO_HIGH_FREQ_MINUTES or 10)) * 60 high_freq_poll_seconds = _clamp_positive_seconds(settings.CHATAPI_ASYNC_VIDEO_HIGH_FREQ_POLL_SECONDS, 30) initial_backoff_seconds = _clamp_positive_seconds(settings.CHATAPI_ASYNC_VIDEO_BACKOFF_INITIAL_SECONDS, 60) multiplier = max(1, int(settings.CHATAPI_ASYNC_VIDEO_BACKOFF_MULTIPLIER or 2)) max_backoff_seconds = _clamp_positive_seconds(settings.CHATAPI_ASYNC_VIDEO_BACKOFF_MAX_SECONDS, 3600) if elapsed_seconds < high_freq_seconds: delay_seconds = high_freq_poll_seconds interval_seconds = int(task.poll_interval_seconds or 0) reason = "video_high_freq_poll" else: previous_interval = int(task.poll_interval_seconds or 0) if previous_interval < initial_backoff_seconds: interval_seconds = initial_backoff_seconds else: interval_seconds = min(previous_interval * multiplier, max_backoff_seconds) delay_seconds = interval_seconds reason = "video_backoff_poll" next_poll_at = current_time + timedelta(seconds=delay_seconds) final_poll = False if deadline_at and next_poll_at >= deadline_at: next_poll_at = deadline_at delay_seconds = max(0, int((deadline_at - current_time).total_seconds())) final_poll = delay_seconds <= 0 reason = "video_schedule_to_final_deadline" direct_max = max(0, int(settings.CHATAPI_ASYNC_VIDEO_DIRECT_COUNTDOWN_MAX_SECONDS or 300)) return PollScheduleDecision( delay_seconds=delay_seconds, next_poll_at=next_poll_at, poll_interval_seconds=interval_seconds, direct_countdown=delay_seconds <= direct_max, final_poll=final_poll, reason=reason, ) def build_default_poll_schedule( task: ChatGenerationTask, *, now: datetime | None = None, delay_seconds: int | None = None, reason: str = "default_poll", ) -> PollScheduleDecision: """图片和旧逻辑兼容用的固定间隔轮询计划。""" current_time = now or utc_now() delay = _clamp_positive_seconds(delay_seconds, int(settings.CHATAPI_ASYNC_POLL_INTERVAL_SECONDS or 30)) next_poll_at = current_time + timedelta(seconds=delay) return PollScheduleDecision( delay_seconds=delay, next_poll_at=next_poll_at, poll_interval_seconds=int(getattr(task, "poll_interval_seconds", 0) or 0), direct_countdown=True, final_poll=False, reason=reason, )