Xiaoxia AI
76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
from worker_app.celery_app import celery_app
|
|
from worker_app.db import SessionLocal
|
|
|
|
from packages.adapters.sqlalchemy_impl.classification_job_repository import (
|
|
SQLAlchemyClassificationJobRepository,
|
|
)
|
|
from packages.domain import (
|
|
AssetClassification,
|
|
ClassificationJob,
|
|
ClassificationJobStatus,
|
|
)
|
|
|
|
|
|
@celery_app.task(name="worker.classify_asset")
|
|
def classify_asset(job_id: str) -> dict:
|
|
"""
|
|
Classify asset task.
|
|
|
|
Steps:
|
|
1. Fetch ClassificationJob from repository
|
|
2. Fetch Asset from repository
|
|
3. Run classification model (placeholder: mock classification)
|
|
4. Update ClassificationJob with result
|
|
5. Return result
|
|
"""
|
|
# 创建数据库 session 和 repository
|
|
session = SessionLocal()
|
|
try:
|
|
job_repo = SQLAlchemyClassificationJobRepository(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()
|
|
|
|
# Mock classification (in real implementation: use ML model, vision API, etc.)
|
|
# For now, randomly classify based on asset_id hash
|
|
asset_id_hash = sum(ord(c) for c in job.asset_id)
|
|
classifications = list(AssetClassification)
|
|
classification = classifications[asset_id_hash % len(classifications)]
|
|
confidence = 0.85
|
|
|
|
# Update job status to COMPLETED
|
|
job.status = ClassificationJobStatus.COMPLETED
|
|
job.classification = classification.value
|
|
job.confidence = confidence
|
|
job_repo.update(job)
|
|
session.commit()
|
|
|
|
return {
|
|
"status": "completed",
|
|
"job_id": job.id,
|
|
"classification": classification.value,
|
|
"confidence": confidence,
|
|
}
|
|
except Exception as e:
|
|
session.rollback()
|
|
# Update job status to FAILED
|
|
job.status = ClassificationJobStatus.FAILED
|
|
job.error_message = str(e)
|
|
job_repo.update(job)
|
|
session.commit()
|
|
|
|
return {
|
|
"status": "failed",
|
|
"job_id": job.id,
|
|
"error": str(e),
|
|
}
|
|
finally:
|
|
session.close()
|