xiaoxia-saas/apps/worker/worker_app/tasks/tts_synthesis.py

"""TTS synthesis tasks - process TTS synthesis requests via CosyVoice API."""

import logging

from celery import Task
from celery.exceptions import Retry
from worker_app.celery_app import celery_app
from worker_app.db import SessionLocal

from packages.adapters.sqlalchemy_impl.tts_job_repository import (
    SQLAlchemyTTSJobRepository,
)
from packages.application.cosyvoice_service import (
    CosyVoiceError,
    CosyVoiceService,
    CosyVoiceTimeoutError,
)
from packages.application.tts_job.workflow import TTSWorkflowService

logger = logging.getLogger(__name__)


@celery_app.task(bind=True, max_retries=3, name="worker.process_tts_synthesis")
def process_tts_synthesis(self: Task, job_id: str) -> dict:
    """处理 TTS 合成任务。

    通过 TTSWorkflowService.poll_and_process_synthesis() 轮询 CosyVoice
    合成任务状态，更新 TTSJob。
    超时自动重试（最多 3 次），其他错误标记 job 为 failed。

    Args:
        job_id: TTSJob ID

    Returns:
        dict: {"ok": True, "job_id": str, "audio_url": str} 或
              {"ok": False, "job_id": str, "error": str}
    """
    session = None
    try:
        session = SessionLocal()
        repo = SQLAlchemyTTSJobRepository(session)
        workflow = TTSWorkflowService(
            repository=repo,
            cosyvoice_service=CosyVoiceService(),
        )

        updated_job = workflow.poll_and_process_synthesis(job_id, timeout=120)
        session.commit()

        logger.info(f"TTS synthesis completed: job_id={job_id}, " f"audio_url={updated_job.output_audio_url}")
        return {
            "ok": True,
            "job_id": job_id,
            "audio_url": updated_job.output_audio_url,
        }

    except Retry:
        # Celery Retry 异常必须向上传播，不能被后续 except 捕获
        raise

    except CosyVoiceTimeoutError as e:
        logger.warning(f"TTS synthesis timeout for {job_id}: {e}")
        if session is not None:
            session.rollback()
        # 超时重试，指数退避
        raise self.retry(exc=e, countdown=30) from e

    except CosyVoiceError as e:
        logger.error(f"TTS synthesis failed for {job_id}: {e}")
        if session is not None:
            session.rollback()
        # 标记 job 为 failed
        try:
            if session is not None:
                job = repo.get(job_id)
                if job is not None:
                    job.mark_failed(str(e))
                    repo.update(job)
                    session.commit()
        except Exception as inner_e:
            logger.error(f"Failed to mark job as failed: {inner_e}")
            if session is not None:
                session.rollback()
        return {"ok": False, "job_id": job_id, "error": str(e)}

    except Exception as e:
        logger.error(f"TTS synthesis unexpected error for {job_id}: {e}")
        if session is not None:
            session.rollback()
        # 标记 job 为 failed
        try:
            if session is not None:
                job = repo.get(job_id)
                if job is not None:
                    job.mark_failed(str(e))
                    repo.update(job)
                    session.commit()
        except Exception as inner_e:
            logger.error(f"Failed to mark job as failed: {inner_e}")
            if session is not None:
                session.rollback()
        return {"ok": False, "job_id": job_id, "error": str(e)}

    finally:
        if session is not None:
            session.close()


@celery_app.task(bind=True, max_retries=2, name="worker.process_tts_segment_synthesis")
def process_tts_segment_synthesis(self: Task, job_id: str) -> dict:
    """分段合成轮询任务 — 轮询多个 CosyVoice 子任务并合并音频。

    与 process_tts_synthesis 类似，但超时更长（300s），
    因为分段任务需要等待所有子任务完成。
    """
    session = None
    try:
        session = SessionLocal()
        repo = SQLAlchemyTTSJobRepository(session)
        workflow = TTSWorkflowService(
            repository=repo,
            cosyvoice_service=CosyVoiceService(),
        )

        updated_job = workflow.poll_and_process_synthesis(job_id, timeout=300)
        session.commit()

        logger.info(f"TTS segment synthesis completed: job_id={job_id}, " f"audio_url={updated_job.output_audio_url}")
        return {
            "ok": True,
            "job_id": job_id,
            "audio_url": updated_job.output_audio_url,
        }

    except Retry:
        raise

    except CosyVoiceTimeoutError as e:
        logger.warning(f"TTS segment synthesis timeout for {job_id}: {e}")
        if session is not None:
            session.rollback()
        raise self.retry(exc=e, countdown=60) from e

    except CosyVoiceError as e:
        logger.error(f"TTS segment synthesis failed for {job_id}: {e}")
        if session is not None:
            session.rollback()
        try:
            if session is not None:
                job = repo.get(job_id)
                if job is not None:
                    job.mark_failed(str(e))
                    repo.update(job)
                    session.commit()
        except Exception as inner_e:
            logger.error(f"Failed to mark job as failed: {inner_e}")
            if session is not None:
                session.rollback()
        return {"ok": False, "job_id": job_id, "error": str(e)}

    except Exception as e:
        logger.error(f"TTS segment synthesis unexpected error for {job_id}: {e}")
        if session is not None:
            session.rollback()
        try:
            if session is not None:
                job = repo.get(job_id)
                if job is not None:
                    job.mark_failed(str(e))
                    repo.update(job)
                    session.commit()
        except Exception as inner_e:
            logger.error(f"Failed to mark job as failed: {inner_e}")
            if session is not None:
                session.rollback()
        return {"ok": False, "job_id": job_id, "error": str(e)}

    finally:
        if session is not None:
            session.close()