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", "自动"
        CHAT = "chat", "聊天补全"
        GEMINI = "gemini", "Gemini"
        IMAGES = "images", "图片生成"
        IMAGES_EDITS = "images_edits", "图片编辑"

    name = models.CharField("名称", max_length=128, unique=True)
    url = models.URLField("接口 URL", max_length=500)
    model = models.CharField("模型标识", max_length=128)
    api_type = models.CharField(
        "API 类型",
        max_length=32,
        choices=ApiType.choices,
        default=ApiType.AUTO,
    )
    api_key_encrypted = models.TextField("加密 API Key", blank=True, editable=False)
    capabilities = models.JSONField("能力", default=list, blank=True)
    timeout_seconds = models.PositiveIntegerField(
        "请求超时秒数",
        default=0,
        help_text="0 表示使用 provider 默认值。",
    )
    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", "生成标题"
        IMAGE = "image", "生成图片"
        VISION = "vision", "图片理解"

    alias = models.SlugField("能力别名", max_length=64)
    operation_type = models.CharField("操作类型", max_length=32, choices=OperationType.choices)
    ai_model = models.ForeignKey(
        AiModel,
        verbose_name="AI 模型",
        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_ALIAS = "model_alias", "能力别名"

    class Action(models.TextChoices):
        CREATE = "create", "创建"
        UPDATE = "update", "更新"
        DELETE = "delete", "删除"

    actor = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        verbose_name="操作人",
        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("目标 ID", 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}"