cmhub/apps/ai/models.py

from __future__ import annotations

from django.conf import settings
from django.core.exceptions import ValidationError
from django.db import models

from apps.ai.providers import ResolvedModel

from .security import decrypt_api_key, encrypt_api_key, mask_secret


class AiModel(models.Model):
    class ApiType(models.TextChoices):
        AUTO = "auto", "Auto"
        CHAT = "chat", "Chat completions"
        GEMINI = "gemini", "Gemini"
        IMAGES = "images", "Images generations"
        IMAGES_EDITS = "images_edits", "Images edits"

    name = models.CharField(max_length=128, unique=True)
    url = models.URLField(max_length=500)
    model = models.CharField(max_length=128)
    api_type = models.CharField(
        max_length=32,
        choices=ApiType.choices,
        default=ApiType.AUTO,
    )
    api_key_encrypted = models.TextField(blank=True, editable=False)
    capabilities = models.JSONField(default=list, blank=True)
    timeout_seconds = models.PositiveIntegerField(
        default=0,
        help_text="0 means use provider resolution defaults.",
    )
    connect_timeout_seconds = models.PositiveIntegerField(default=30)
    extra_body = models.JSONField(default=dict, blank=True)
    is_active = models.BooleanField(default=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        db_table = "ai_model"
        verbose_name = "AI 模型"
        verbose_name_plural = "AI 模型"
        ordering = ("name",)

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

    def set_api_key(self, raw_api_key: str) -> None:
        self.api_key_encrypted = encrypt_api_key(raw_api_key)

    def get_api_key(self) -> str:
        return decrypt_api_key(self.api_key_encrypted)

    @property
    def has_api_key(self) -> bool:
        return bool(self.api_key_encrypted)

    @property
    def api_key_masked(self) -> str:
        return mask_secret(self.api_key_encrypted)

    def capabilities_set(self) -> frozenset[str]:
        raw = self.capabilities or []
        if isinstance(raw, dict):
            return frozenset(
                str(key)
                for key, value in raw.items()
                if key in {"text", "image", "vision"} and bool(value)
            )
        if isinstance(raw, str):
            return frozenset(item.strip() for item in raw.split(",") if item.strip())
        return frozenset(str(item) for item in raw)

    def to_resolved_model(self) -> ResolvedModel:
        return ResolvedModel(
            name=self.name,
            url=self.url,
            model=self.model,
            api_key=self.get_api_key(),
            api_type=self.api_type,
            timeout_seconds=self.timeout_seconds,
            connect_timeout_seconds=self.connect_timeout_seconds,
            extra_body=dict(self.extra_body or {}),
            capabilities=self.capabilities_set(),
        )


class ModelAlias(models.Model):
    class OperationType(models.TextChoices):
        TITLE = "title", "Generate title"
        IMAGE = "image", "Generate image"

    alias = models.SlugField(max_length=64)
    operation_type = models.CharField(max_length=32, choices=OperationType.choices)
    ai_model = models.ForeignKey(
        AiModel,
        on_delete=models.PROTECT,
        related_name="aliases",
    )
    is_default = models.BooleanField(default=False)
    is_active = models.BooleanField(default=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        db_table = "model_alias"
        verbose_name = "能力别名"
        verbose_name_plural = "能力别名"
        ordering = ("operation_type", "alias")
        constraints = [
            models.UniqueConstraint(
                fields=("operation_type", "alias"),
                name="unique_model_alias_per_operation",
            ),
        ]
        indexes = [
            models.Index(fields=("operation_type", "is_default")),
            models.Index(fields=("operation_type", "is_active")),
        ]

    def __str__(self) -> str:
        marker = " default" if self.is_default else ""
        return f"{self.operation_type}:{self.alias}{marker}"

    def clean(self) -> None:
        super().clean()
        if self.is_default:
            duplicate_default = ModelAlias.objects.filter(
                operation_type=self.operation_type,
                is_default=True,
            ).exclude(pk=self.pk)
            if duplicate_default.exists():
                raise ValidationError(
                    {"is_default": "Only one default alias is allowed per operation."}
                )

    def save(self, *args, **kwargs) -> None:
        self.full_clean()
        super().save(*args, **kwargs)


class AiConfigAuditLog(models.Model):
    class TargetType(models.TextChoices):
        AI_MODEL = "ai_model", "AI model"
        MODEL_ALIAS = "model_alias", "Model alias"

    class Action(models.TextChoices):
        CREATE = "create", "Create"
        UPDATE = "update", "Update"
        DELETE = "delete", "Delete"

    actor = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="+",
    )
    action = models.CharField(max_length=16, choices=Action.choices)
    target_type = models.CharField(max_length=32, choices=TargetType.choices)
    target_id = models.PositiveBigIntegerField(null=True, blank=True)
    target_repr = models.CharField(max_length=255)
    changed_fields = models.JSONField(default=list, blank=True)
    changes = models.JSONField(default=dict, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)

    class Meta:
        db_table = "ai_config_audit_log"
        verbose_name = "AI 配置审计日志"
        verbose_name_plural = "AI 配置审计日志"
        ordering = ("-created_at", "-id")
        indexes = [
            models.Index(fields=("target_type", "target_id")),
            models.Index(fields=("action", "created_at")),
            models.Index(fields=("actor", "created_at")),
        ]

    def __str__(self) -> str:
        return f"{self.created_at:%Y-%m-%d %H:%M:%S} {self.action} {self.target_repr}"