xiaoxia-saas/apps/worker/tasks.py

from datetime import datetime, timezone
import random

from app.config import get_settings
from .celery_app import celery_app
from packages.adapters.sqlalchemy_impl.session import SessionLocal, build_session_factory
from packages.adapters.sqlalchemy_impl.ingest_job_repository import SQLAlchemyIngestJobRepository
from packages.adapters.sqlalchemy_impl.asset_repository import SQLAlchemyAssetRepository
from packages.adapters.sqlalchemy_impl.classification_job_repository import SQLAlchemyClassificationJobRepository
from packages.domain import Asset, AssetClassification, ClassificationJob, ClassificationJobStatus, IngestJobStatus

settings = get_settings()
if SessionLocal is None:
    build_session_factory(settings.database_url)


@celery_app.task(name="worker.healthcheck")
def healthcheck() -> dict:
    return {"ok": True, "service": "worker"}


@celery_app.task(name="worker.ingest_asset")
def ingest_asset(job_id: str) -> dict:
    session = SessionLocal()
    try:
        ingest_repo = SQLAlchemyIngestJobRepository(session)
        asset_repo = SQLAlchemyAssetRepository(session)
        classification_repo = SQLAlchemyClassificationJobRepository(session)

        job = ingest_repo.get(job_id)
        if job is None:
            return {"ok": False, "error": f"job {job_id} not found"}

        job.status = IngestJobStatus.PROCESSING
        job.updated_at = datetime.now(timezone.utc)
        ingest_repo.update(job)

        storage_key = job.storage_key
        filename = storage_key.split("/")[-1]
        lower_name = filename.lower()
        if lower_name.endswith((".mp4", ".mov", ".avi", ".mkv")):
            mime_type = "video/mp4"
        elif lower_name.endswith((".mp3", ".wav", ".aac")):
            mime_type = "audio/mpeg"
        elif lower_name.endswith((".jpg", ".jpeg")):
            mime_type = "image/jpeg"
        elif lower_name.endswith((".png", ".webp", ".gif")):
            mime_type = "image/png"
        else:
            mime_type = "application/octet-stream"

        asset = Asset.create(
            workspace_id=job.workspace_id,
            project_id=job.project_id,
            library_id=job.library_id,
            name=filename,
            storage_key=storage_key,
            mime_type=mime_type,
            metadata={"source": "ingest_task", "auto_classification": "queued"},
        )
        asset_repo.create(asset)

        job.status = IngestJobStatus.COMPLETED
        job.result_asset_id = asset.id
        job.updated_at = datetime.now(timezone.utc)
        ingest_repo.update(job)

        classification_job = ClassificationJob.create(
            workspace_id=job.workspace_id,
            project_id=job.project_id,
            asset_id=asset.id,
        )
        classification_repo.create(classification_job)
        celery_app.send_task("worker.classify_asset", args=[classification_job.id])

        return {
            "ok": True,
            "job_id": job.id,
            "asset_id": asset.id,
            "classification_job_id": classification_job.id,
        }
    except Exception as error:
        try:
            ingest_repo = SQLAlchemyIngestJobRepository(session)
            job = ingest_repo.get(job_id)
            if job is not None:
                job.status = IngestJobStatus.FAILED
                job.error_message = str(error)
                job.updated_at = datetime.now(timezone.utc)
                ingest_repo.update(job)
        except Exception:
            pass
        return {"ok": False, "job_id": job_id, "error": str(error)}
    finally:
        session.close()


@celery_app.task(name="worker.classify_asset")
def classify_asset(job_id: str) -> dict:
    session = SessionLocal()
    try:
        classification_repo = SQLAlchemyClassificationJobRepository(session)
        asset_repo = SQLAlchemyAssetRepository(session)

        job = classification_repo.get(job_id)
        if job is None:
            return {"ok": False, "error": f"classification job {job_id} not found"}

        job.status = ClassificationJobStatus.PROCESSING
        job.updated_at = datetime.now(timezone.utc)
        classification_repo.update(job)

        asset = asset_repo.get(job.asset_id)
        if asset is None:
            raise ValueError(f"asset {job.asset_id} not found")

        name = asset.name.lower()
        if any(token in name for token in ["food", "meal", "cook"]):
            classification = AssetClassification.FOOD.value
        elif any(token in name for token in ["person", "human", "portrait"]):
            classification = AssetClassification.PERSON.value
        elif any(token in name for token in ["music", "song", "audio"]):
            classification = AssetClassification.MUSIC.value
        elif any(token in name for token in ["product", "sku", "item"]):
            classification = AssetClassification.PRODUCT.value
        elif any(token in name for token in ["animal", "pet", "cat", "dog"]):
            classification = AssetClassification.ANIMAL.value
        elif any(token in name for token in ["sport", "run", "ball"]):
            classification = AssetClassification.SPORT.value
        elif any(token in name for token in ["tech", "phone", "device", "pc"]):
            classification = AssetClassification.TECH.value
        elif any(token in name for token in ["view", "travel", "mountain", "sea"]):
            classification = AssetClassification.SCENIC.value
        else:
            classification = AssetClassification.OTHER.value

        confidence = round(random.uniform(0.72, 0.96), 2)
        asset.metadata = {
            **asset.metadata,
            "classification": classification,
            "classification_confidence": confidence,
            "auto_classification": "completed",
        }
        asset_repo.update(asset)

        job.status = ClassificationJobStatus.COMPLETED
        job.classification = classification
        job.confidence = confidence
        job.updated_at = datetime.now(timezone.utc)
        classification_repo.update(job)

        return {
            "ok": True,
            "job_id": job.id,
            "asset_id": asset.id,
            "classification": classification,
        }
    except Exception as error:
        try:
            classification_repo = SQLAlchemyClassificationJobRepository(session)
            job = classification_repo.get(job_id)
            if job is not None:
                job.status = ClassificationJobStatus.FAILED
                job.error_message = str(error)
                job.updated_at = datetime.now(timezone.utc)
                classification_repo.update(job)
        except Exception:
            pass
        return {"ok": False, "job_id": job_id, "error": str(error)}
    finally:
        session.close()