celery_app.py

import os
from celery import Celery
from database import SessionLocal
from models import Submission, ReviewConfig, ReviewResult
from agent import build_review_messages, generate_review
from logger import get_logger

logger = get_logger("celery_tasks")

redis_url = os.getenv("CELERY_BROKER_URL", "redis://localhost:6379/0")

celery = Celery(
    "labelhub_tasks",
    broker=redis_url,
    backend=redis_url
)

celery.conf.update(
    task_serializer="json",
    accept_content=["json"],
    result_serializer="json",
    timezone="Asia/Shanghai",
    enable_utc=True,
)

@celery.task(bind=True, max_retries=3, default_retry_delay=5)
def process_submission(self, submission_id_db: int):
    """
    处理标注数据审核
    幂等性：如果任务状态不是 pending，则直接返回
    """
    db = SessionLocal()
    submission = None
    try:
        # 获取提交记录
        submission = db.query(Submission).filter(Submission.id == submission_id_db).first()
        if not submission:
            logger.error(f"Submission {submission_id_db} not found")
            return {"status": "error", "message": "Submission not found"}

        # 幂等性校验：只有 pending 的才处理
        if submission.status != "pending":
            logger.info(f"Submission {submission.submission_id} already processed or in progress. Status: {submission.status}")
            return {"status": "skipped", "message": "Submission already processed or in progress"}

        logger.info(f"Processing submission: {submission.submission_id}")
        # 更新状态为 processing
        submission.status = "processing"
        db.commit()

        # 获取配置
        config = db.query(ReviewConfig).filter(ReviewConfig.task_name == submission.task_name).first()
        if not config:
            submission.status = "failed"
            db.commit()
            return {"status": "error", "message": "ReviewConfig not found"}

        # 调用 Agent
        messages = build_review_messages(config.prompt_template, submission.data, submission.data_type)
        
        try:
            logger.info(f"Calling LLM ({config.ai_model}) for submission {submission.submission_id} with dimensions: {config.dimensions}")
            result_data = generate_review(messages, config.dimensions, ai_model=config.ai_model)
            logger.info(f"LLM review result for {submission.submission_id}: {result_data}")
        except Exception as exc:
            # 失败重试
            logger.warning(f"LLM review failed for {submission.submission_id}, retrying: {str(exc)}")
            submission.status = "pending" # 恢复状态以便重试
            db.commit()
            raise self.retry(exc=exc)

        # 记录审核结果
        review_result = ReviewResult(
            submission_id=submission.submission_id,
            scores=result_data["scores"],
            ai_comment=result_data["ai_comment"],
            conclusion=result_data["conclusion"],
            original_prompt=config.prompt_template if submission.data_type == "image_url" else str(messages)
        )
        db.add(review_result)
        
        submission.status = "reviewed"
        db.commit()
        
        return {"status": "success", "submission_id": submission.submission_id}
        
    except Exception as e:
        db.rollback()
        # 如果不是 Retry 异常，则设为 failed
        if submission and type(e).__name__ != 'Retry':
            submission.status = "failed"
            db.commit()
        raise e
    finally:
        db.close()