feat: add ai model alias configuration

2026-07-02 11:07:44 +08:00
parent c9fbabb04f
commit 18ac054dc5
23 changed files with 955 additions and 48 deletions
@@ -1,3 +1,135 @@
+from __future__ import annotations
+
+from django.core.exceptions import ValidationError
 from django.db import models

-# Create your models here.
+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"
+        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"
+        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)