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 # 把分类结果写入 metadata,供列表筛选和智能视图使用 asset.metadata = { **(asset.metadata or {}), "classification": classification, "classification_confidence": confidence, } 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()