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()