cmhub/apps/api/admin.py

from django.contrib import admin
from django.db.models import Count, Q

from .models import ImageGenerationTask, ImageGenerationTaskInput


class ImageGenerationTaskInputInline(admin.TabularInline):
    model = ImageGenerationTaskInput
    extra = 0
    can_delete = False
    fields = ("ordinal", "image", "mime_type", "filename", "created_at")
    readonly_fields = fields


class ImageInputTypeFilter(admin.SimpleListFilter):
    title = "输入图片类型"
    parameter_name = "input_image_type"

    def lookups(self, request, model_admin):
        return (
            ("single", "单图生图"),
            ("multiple", "多图生图"),
        )

    def queryset(self, request, queryset):
        value = self.value()
        if value not in {"single", "multiple"}:
            return queryset

        input_counts = queryset.annotate(input_image_count=Count("input_images"))
        if value == "multiple":
            task_ids = input_counts.filter(input_image_count__gte=2).values("pk")
            return queryset.filter(pk__in=task_ids)

        new_single_ids = input_counts.filter(input_image_count=1).values("pk")
        legacy_single_ids = queryset.filter(input_images__isnull=True).exclude(input_image="").values("pk")
        return queryset.filter(Q(pk__in=new_single_ids) | Q(pk__in=legacy_single_ids))


@admin.register(ImageGenerationTask)
class ImageGenerationTaskAdmin(admin.ModelAdmin):
    change_form_template = "admin/api/imagegenerationtask/change_form.html"
    inlines = (ImageGenerationTaskInputInline,)
    list_display = (
        "task_id",
        "user",
        "status",
        "attempt_count",
        "next_attempt_at",
        "points_balance_after_charge",
        "created_at",
        "finished_at",
    )
    list_filter = (ImageInputTypeFilter, "status", "created_at", "next_attempt_at", "finished_at")
    search_fields = (
        "task_id",
        "user__username",
        "api_key__key_prefix",
        "call_record__id",
        "idempotency_key",
        "error_code",
    )
    readonly_fields = (
        "task_id",
        "user",
        "api_key",
        "call_record",
        "idempotency_key_hash",
        "request_hash",
        "request_payload",
        "input_image",
        "result_url",
        "points_balance_after_charge",
        "started_at",
        "finished_at",
        "expires_at",
        "next_attempt_at",
        "locked_at",
        "lease_expires_at",
        "heartbeat_at",
        "worker_id",
        "attempt_count",
        "created_at",
        "updated_at",
    )