yingqing/core/backend/apps/ai/models.py

from django.db import models

from apps.common.models import TeamOwnedModel, TimeStampedModel


class ModelProvider(TimeStampedModel):
    class Status(models.TextChoices):
        ACTIVE = "active", "Active"
        DISABLED = "disabled", "Disabled"

    name = models.CharField(max_length=64, unique=True)
    display_name = models.CharField(max_length=128)
    status = models.CharField(max_length=24, choices=Status.choices, default=Status.ACTIVE)
    base_url = models.URLField(blank=True)
    # 站级 API Key(中转站)。可插拔:换站 = 改这一行的 base_url + api_key,零改代码。
    # 留空则 services 层按 provider.name 回退 settings.PROVIDER_KEYS(.env),避免密钥写死/进库。
    api_key = models.CharField(max_length=255, blank=True)
    metadata = models.JSONField(default=dict, blank=True)

    def __str__(self) -> str:
        return self.display_name


class ModelConfig(TimeStampedModel):
    class Capability(models.TextChoices):
        TEXT = "text", "Text"
        IMAGE = "image", "Image"
        VIDEO = "video", "Video"
        VISION = "vision", "Vision"
        AUDIO = "audio", "Audio"
        EXPORT = "export", "Export"

    class Status(models.TextChoices):
        ACTIVE = "active", "Active"
        DISABLED = "disabled", "Disabled"

    provider = models.ForeignKey(ModelProvider, on_delete=models.CASCADE, related_name="models")
    name = models.CharField(max_length=128)
    display_name = models.CharField(max_length=128)
    capability = models.CharField(max_length=32, choices=Capability.choices)
    endpoint = models.CharField(max_length=255, blank=True)
    unit_price = models.DecimalField(max_digits=12, decimal_places=4, default=0)
    status = models.CharField(max_length=24, choices=Status.choices, default=Status.ACTIVE)
    rate_limit_per_minute = models.PositiveIntegerField(default=60)
    # 平台超管显式钦定的「该能力默认模型」。get_default_model 优先取它,未设则回落最早 active(零回归)。
    is_default = models.BooleanField(default=False)
    metadata = models.JSONField(default=dict, blank=True)

    class Meta:
        unique_together = [("provider", "name", "capability")]

    def __str__(self) -> str:
        return f"{self.provider.name}:{self.name}:{self.capability}"


class AITask(TeamOwnedModel):
    class Type(models.TextChoices):
        SCRIPT_GENERATION = "script_generation", "Script Generation"
        SCRIPT_OPTIMIZATION = "script_optimization", "Script Optimization"
        ENTITY_EXTRACTION = "entity_extraction", "Entity Extraction"
        PRODUCT_IMAGE = "product_image", "Product Image"
        PERSON_IMAGE = "person_image", "Person Image"
        SCENE_IMAGE = "scene_image", "Scene Image"
        STORYBOARD = "storyboard", "Storyboard"
        VIDEO_SEGMENT = "video_segment", "Video Segment"
        VOICEOVER = "voiceover", "Voiceover"
        EXPORT = "export", "Export"

    class Status(models.TextChoices):
        CREATED = "created", "Created"
        RESERVED = "reserved", "Reserved"
        SUBMITTED = "submitted", "Submitted"
        POLLING = "polling", "Polling"
        POSTPROCESSING = "postprocessing", "Postprocessing"
        SUCCEEDED = "succeeded", "Succeeded"
        FAILED = "failed", "Failed"
        COMPENSATING = "compensating", "Compensating"
        CANCELLED = "cancelled", "Cancelled"

    project = models.ForeignKey(
        "projects.Project",
        on_delete=models.CASCADE,
        null=True,
        blank=True,
        related_name="ai_tasks",
    )
    task_type = models.CharField(max_length=48, choices=Type.choices)
    status = models.CharField(max_length=32, choices=Status.choices, default=Status.CREATED)
    model_config = models.ForeignKey(ModelConfig, on_delete=models.PROTECT, related_name="tasks")
    idempotency_key = models.CharField(max_length=128, unique=True)
    provider_task_id = models.CharField(max_length=255, blank=True)
    request_payload = models.JSONField(default=dict, blank=True)
    response_payload = models.JSONField(default=dict, blank=True)
    estimated_cost = models.DecimalField(max_digits=12, decimal_places=4, default=0)
    actual_cost = models.DecimalField(max_digits=12, decimal_places=4, default=0)
    error_code = models.CharField(max_length=64, blank=True)
    error_message = models.TextField(blank=True)
    submitted_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)

    class Meta:
        indexes = [
            models.Index(fields=["team", "status"]),
            models.Index(fields=["project", "task_type"]),
            models.Index(fields=["provider_task_id"]),
            # 任务历史默认按团队 + 创建时间倒序(AI 工具页 / asset-factory)
            models.Index(fields=["team", "-created_at"]),
        ]

    def __str__(self) -> str:
        return f"{self.task_type}:{self.status}:{self.id}"


class QualityWord(TimeStampedModel):
    """平台单层质量词配置:按生成阶段(stage)+ 槽位(slot)挂若干质量词,
    生成侧拼提示词时优先读取本配置;**无配置则回落各 builder 的写死值**(零回归)。
    全平台共用一套(不挂 team),团队级覆盖以后再说。"""

    class Stage(models.TextChoices):
        PERSON = "person", "人物立绘"
        MODEL_TRYON = "model_tryon", "模特上身图"
        STORYBOARD = "storyboard", "故事板"
        VIDEO = "video", "视频"

    stage = models.CharField(max_length=24, choices=Stage.choices)
    slot = models.CharField(max_length=24, default="quality")
    text = models.CharField(max_length=200)
    sort = models.IntegerField(default=0)
    enabled = models.BooleanField(default=True)

    class Meta:
        ordering = ["stage", "sort", "created_at"]
        indexes = [models.Index(fields=["stage", "slot"])]

    def __str__(self) -> str:
        return f"{self.stage}/{self.slot}/{self.text}"