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

from celery.utils.log import get_task_logger
from worker_app.celery_app import celery_app

from packages.adapters.sqlalchemy_impl.asset_repository import (
    SQLAlchemyAssetRepository,
)
from packages.adapters.sqlalchemy_impl.classification_job_repository import (
    SQLAlchemyClassificationJobRepository,
)
from packages.domain import (
    ClassificationJobStatus,
    ClassificationStatus,
)

from .asset_analyzer import classify_asset_real

logger = get_task_logger(__name__)


@celery_app.task(bind=True, name="worker.classify_asset", max_retries=2)
def classify_asset(self, job_id: str) -> dict:
    """
    Classify asset task.

    Steps:
    1. Fetch ClassificationJob from repository
    2. Fetch Asset from repository
    3. Run real classification based on video features
    4. Update ClassificationJob with result
    5. Update Asset with classification and status
    6. Return result
    """
    from worker_app.db import SessionLocal

    # 创建数据库 session 和 repository
    session = SessionLocal()
    try:
        job_repo = SQLAlchemyClassificationJobRepository(session)
        asset_repo = SQLAlchemyAssetRepository(session)

        job = job_repo.get(job_id)
        if job is None:
            return {"status": "failed", "error": "job not found"}

        try:
            # Update job status to PROCESSING
            job.status = ClassificationJobStatus.PROCESSING
            job_repo.update(job)
            session.commit()

            # Get the asset to find the video path
            asset = asset_repo.get(job.asset_id)
            if asset is None:
                raise ValueError(f"Asset not found: {job.asset_id}")

            # Determine media path from storage_key
            # In production, this would be a full URL/path to the media file
            video_path = asset.storage_key

            # Run real classification
            classification, confidence = classify_asset_real(video_path)

            # Update job with classification result
            job.status = ClassificationJobStatus.COMPLETED
            job.classification = classification
            job.confidence = confidence
            job_repo.update(job)

            # Update asset with classification status and result
            asset.classification_status = ClassificationStatus.COMPLETED
            asset_repo.update(asset)

            session.commit()

            logger.info(
                f"Classification completed for asset {asset.id}: " f"category={classification}, confidence={confidence}"
            )

            return {
                "status": "completed",
                "job_id": job.id,
                "classification": classification,
                "confidence": confidence,
            }
        except Exception as e:
            session.rollback()
            logger.error(f"Classification failed for job {job_id}: {e}")

            # Update job status to FAILED
            job.status = ClassificationJobStatus.FAILED
            job.error_message = str(e)
            job_repo.update(job)

            # Update asset classification status to FAILED
            asset = asset_repo.get(job.asset_id)
            if asset:
                asset.classification_status = ClassificationStatus.FAILED
                asset_repo.update(asset)

            session.commit()

            return {
                "status": "failed",
                "job_id": job.id,
                "error": str(e),
            }
    finally:
        session.close()