136 lines
4.5 KiB
Python
136 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
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"
|
|
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)
|