feat: add ai model alias configuration
This commit is contained in:
QiuSW
+113
-1
@@ -1,3 +1,115 @@
|
||||
from django import forms
|
||||
from django.contrib import admin
|
||||
|
||||
# Register your models here.
|
||||
from .models import AiModel, ModelAlias
|
||||
from .security import AiKeyEncryptionError, encrypt_api_key
|
||||
|
||||
|
||||
class AiModelAdminForm(forms.ModelForm):
|
||||
api_key = forms.CharField(
|
||||
label="API key",
|
||||
required=False,
|
||||
widget=forms.PasswordInput(render_value=False),
|
||||
help_text="Leave blank to keep the existing encrypted key.",
|
||||
)
|
||||
|
||||
class Meta:
|
||||
model = AiModel
|
||||
fields = (
|
||||
"name",
|
||||
"url",
|
||||
"model",
|
||||
"api_type",
|
||||
"api_key",
|
||||
"capabilities",
|
||||
"timeout_seconds",
|
||||
"connect_timeout_seconds",
|
||||
"extra_body",
|
||||
"is_active",
|
||||
)
|
||||
|
||||
def clean(self):
|
||||
cleaned_data = super().clean()
|
||||
api_key = cleaned_data.get("api_key")
|
||||
if not self.instance.pk and not api_key:
|
||||
raise forms.ValidationError("API key is required when creating an AI model.")
|
||||
if api_key:
|
||||
try:
|
||||
self._api_key_encrypted = encrypt_api_key(api_key)
|
||||
except AiKeyEncryptionError as exc:
|
||||
raise forms.ValidationError({"api_key": str(exc)}) from exc
|
||||
else:
|
||||
self._api_key_encrypted = ""
|
||||
return cleaned_data
|
||||
|
||||
def save(self, commit=True):
|
||||
instance = super().save(commit=False)
|
||||
if self._api_key_encrypted:
|
||||
instance.api_key_encrypted = self._api_key_encrypted
|
||||
if commit:
|
||||
instance.save()
|
||||
self.save_m2m()
|
||||
return instance
|
||||
|
||||
|
||||
@admin.register(AiModel)
|
||||
class AiModelAdmin(admin.ModelAdmin):
|
||||
form = AiModelAdminForm
|
||||
list_display = (
|
||||
"name",
|
||||
"model",
|
||||
"api_type",
|
||||
"capabilities_display",
|
||||
"api_key_status",
|
||||
"is_active",
|
||||
"updated_at",
|
||||
)
|
||||
list_filter = ("api_type", "is_active")
|
||||
search_fields = ("name", "model", "url")
|
||||
readonly_fields = ("api_key_status", "created_at", "updated_at")
|
||||
fieldsets = (
|
||||
(None, {"fields": ("name", "url", "model", "api_type", "is_active")}),
|
||||
(
|
||||
"Credentials",
|
||||
{
|
||||
"fields": ("api_key", "api_key_status"),
|
||||
"description": "The stored key is encrypted and never displayed.",
|
||||
},
|
||||
),
|
||||
(
|
||||
"Capabilities and request defaults",
|
||||
{
|
||||
"fields": (
|
||||
"capabilities",
|
||||
"timeout_seconds",
|
||||
"connect_timeout_seconds",
|
||||
"extra_body",
|
||||
)
|
||||
},
|
||||
),
|
||||
("Timestamps", {"fields": ("created_at", "updated_at")}),
|
||||
)
|
||||
|
||||
@admin.display(description="capabilities")
|
||||
def capabilities_display(self, obj):
|
||||
return ", ".join(sorted(obj.capabilities_set()))
|
||||
|
||||
@admin.display(description="API key")
|
||||
def api_key_status(self, obj):
|
||||
return obj.api_key_masked or "not set"
|
||||
|
||||
|
||||
@admin.register(ModelAlias)
|
||||
class ModelAliasAdmin(admin.ModelAdmin):
|
||||
list_display = (
|
||||
"operation_type",
|
||||
"alias",
|
||||
"ai_model",
|
||||
"is_default",
|
||||
"is_active",
|
||||
"updated_at",
|
||||
)
|
||||
list_filter = ("operation_type", "is_default", "is_active")
|
||||
search_fields = ("alias", "ai_model__name", "ai_model__model")
|
||||
autocomplete_fields = ("ai_model",)
|
||||
readonly_fields = ("created_at", "updated_at")
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from apps.ai.providers import ResolvedModel
|
||||
|
||||
from .models import AiModel, ModelAlias
|
||||
|
||||
|
||||
class AliasResolutionError(RuntimeError):
|
||||
"""Base error for alias resolution failures."""
|
||||
|
||||
|
||||
class AliasNotFoundError(AliasResolutionError):
|
||||
"""Raised when an active alias cannot be found."""
|
||||
|
||||
|
||||
class ModelCapabilityError(AliasResolutionError):
|
||||
"""Raised when an alias points to a model without the required capability."""
|
||||
|
||||
|
||||
REQUIRED_CAPABILITIES = {
|
||||
ModelAlias.OperationType.TITLE: "text",
|
||||
ModelAlias.OperationType.IMAGE: "image",
|
||||
}
|
||||
|
||||
|
||||
def resolve_alias(operation_type: str, alias: str | None = None) -> ResolvedModel:
|
||||
"""Resolve an external capability alias to a provider-ready model config."""
|
||||
required_capability = REQUIRED_CAPABILITIES.get(operation_type)
|
||||
if required_capability is None:
|
||||
raise AliasResolutionError(f"unsupported operation_type: {operation_type}")
|
||||
|
||||
queryset = ModelAlias.objects.select_related("ai_model").filter(
|
||||
operation_type=operation_type,
|
||||
is_active=True,
|
||||
ai_model__is_active=True,
|
||||
)
|
||||
if alias:
|
||||
queryset = queryset.filter(alias=alias)
|
||||
else:
|
||||
queryset = queryset.filter(is_default=True)
|
||||
|
||||
model_alias = queryset.order_by("id").first()
|
||||
if model_alias is None:
|
||||
if alias:
|
||||
raise AliasNotFoundError(f"alias not found: {operation_type}:{alias}")
|
||||
raise AliasNotFoundError(f"default alias not found: {operation_type}")
|
||||
|
||||
ai_model: AiModel = model_alias.ai_model
|
||||
capabilities = ai_model.capabilities_set()
|
||||
if required_capability not in capabilities:
|
||||
raise ModelCapabilityError(
|
||||
f"alias {model_alias.alias} maps to model {ai_model.name} without "
|
||||
f"{required_capability} capability"
|
||||
)
|
||||
return ai_model.to_resolved_model()
|
||||
@@ -0,0 +1,158 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Mapping
|
||||
|
||||
from apps.ai.providers.utils import (
|
||||
API_CHAT,
|
||||
API_GEMINI,
|
||||
API_IMAGES,
|
||||
API_IMAGES_EDITS,
|
||||
detect_api_type,
|
||||
)
|
||||
|
||||
from .models import AiModel, ModelAlias
|
||||
|
||||
|
||||
def import_ai_models_config(
|
||||
config: Mapping[str, Any],
|
||||
*,
|
||||
create_default_aliases: bool = False,
|
||||
) -> dict[str, int]:
|
||||
"""Import cmbot-style ai_models.json data into AiModel records."""
|
||||
raw_models = config.get("models")
|
||||
if not isinstance(raw_models, list):
|
||||
raise ValueError("ai_models config must contain a models list")
|
||||
|
||||
created = 0
|
||||
updated = 0
|
||||
imported_models: list[AiModel] = []
|
||||
for raw_model in raw_models:
|
||||
if not isinstance(raw_model, Mapping):
|
||||
raise ValueError("each model config must be an object")
|
||||
ai_model, was_created = _upsert_ai_model(raw_model)
|
||||
imported_models.append(ai_model)
|
||||
if was_created:
|
||||
created += 1
|
||||
else:
|
||||
updated += 1
|
||||
|
||||
aliases = 0
|
||||
if create_default_aliases:
|
||||
aliases = _ensure_default_aliases(imported_models)
|
||||
|
||||
return {"created": created, "updated": updated, "aliases": aliases}
|
||||
|
||||
|
||||
def _upsert_ai_model(raw_model: Mapping[str, Any]) -> tuple[AiModel, bool]:
|
||||
name = str(raw_model.get("name") or raw_model.get("model") or "").strip()
|
||||
if not name:
|
||||
raise ValueError("model config is missing name/model")
|
||||
|
||||
defaults = {
|
||||
"url": str(raw_model.get("url") or ""),
|
||||
"model": str(raw_model.get("model") or ""),
|
||||
"api_type": str(raw_model.get("api_type") or AiModel.ApiType.AUTO),
|
||||
"timeout_seconds": _to_non_negative_int(raw_model.get("timeout_seconds"), 0),
|
||||
"connect_timeout_seconds": _to_positive_int(
|
||||
raw_model.get("connect_timeout_seconds"), 30
|
||||
),
|
||||
"extra_body": dict(raw_model.get("extra_body") or {}),
|
||||
"capabilities": _infer_capabilities(raw_model),
|
||||
"is_active": True,
|
||||
}
|
||||
ai_model, created = AiModel.objects.update_or_create(name=name, defaults=defaults)
|
||||
|
||||
api_key = str(raw_model.get("api_key") or "").strip()
|
||||
if api_key:
|
||||
ai_model.set_api_key(api_key)
|
||||
ai_model.save(update_fields=("api_key_encrypted", "updated_at"))
|
||||
return ai_model, created
|
||||
|
||||
|
||||
def _infer_capabilities(raw_model: Mapping[str, Any]) -> list[str]:
|
||||
explicit = raw_model.get("capabilities")
|
||||
if explicit:
|
||||
if isinstance(explicit, dict):
|
||||
return [
|
||||
str(key)
|
||||
for key, value in explicit.items()
|
||||
if key in {"text", "image", "vision"} and bool(value)
|
||||
]
|
||||
if isinstance(explicit, str):
|
||||
return [item.strip() for item in explicit.split(",") if item.strip()]
|
||||
return [str(item) for item in explicit]
|
||||
|
||||
api_type = str(raw_model.get("api_type") or AiModel.ApiType.AUTO)
|
||||
url = str(raw_model.get("url") or "")
|
||||
resolved_api_type = detect_api_type(url, api_type)
|
||||
name_model_text = f"{raw_model.get('name', '')} {raw_model.get('model', '')}".lower()
|
||||
|
||||
if resolved_api_type == API_IMAGES_EDITS:
|
||||
return ["image", "vision"]
|
||||
if resolved_api_type == API_IMAGES:
|
||||
return ["image"]
|
||||
if resolved_api_type == API_GEMINI:
|
||||
return ["text", "image", "vision"]
|
||||
if resolved_api_type == API_CHAT and (
|
||||
"image" in name_model_text or "banana" in name_model_text
|
||||
):
|
||||
return ["image", "vision"]
|
||||
return ["text", "vision"]
|
||||
|
||||
|
||||
def _ensure_default_aliases(models: list[AiModel]) -> int:
|
||||
aliases = 0
|
||||
text_model = _first_model_with_capability(models, "text")
|
||||
image_model = _first_model_with_capability(models, "image")
|
||||
|
||||
if text_model:
|
||||
ModelAlias.objects.filter(
|
||||
operation_type=ModelAlias.OperationType.TITLE,
|
||||
is_default=True,
|
||||
).update(is_default=False)
|
||||
ModelAlias.objects.update_or_create(
|
||||
operation_type=ModelAlias.OperationType.TITLE,
|
||||
alias="title-standard",
|
||||
defaults={"ai_model": text_model, "is_default": True, "is_active": True},
|
||||
)
|
||||
aliases += 1
|
||||
|
||||
if image_model:
|
||||
ModelAlias.objects.filter(
|
||||
operation_type=ModelAlias.OperationType.IMAGE,
|
||||
is_default=True,
|
||||
).update(is_default=False)
|
||||
ModelAlias.objects.update_or_create(
|
||||
operation_type=ModelAlias.OperationType.IMAGE,
|
||||
alias="image-standard",
|
||||
defaults={"ai_model": image_model, "is_default": True, "is_active": True},
|
||||
)
|
||||
aliases += 1
|
||||
|
||||
return aliases
|
||||
|
||||
|
||||
def _first_model_with_capability(models: list[AiModel], capability: str) -> AiModel | None:
|
||||
for model in models:
|
||||
if capability in model.capabilities_set():
|
||||
return model
|
||||
return None
|
||||
|
||||
|
||||
def _to_non_negative_int(value: Any, default: int) -> int:
|
||||
parsed = _to_int(value, default)
|
||||
return parsed if parsed >= 0 else default
|
||||
|
||||
|
||||
def _to_positive_int(value: Any, default: int) -> int:
|
||||
parsed = _to_int(value, default)
|
||||
return parsed if parsed > 0 else default
|
||||
|
||||
|
||||
def _to_int(value: Any, default: int) -> int:
|
||||
if value is None or value == "":
|
||||
return default
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from django.core.management.base import BaseCommand, CommandError
|
||||
|
||||
from apps.ai.importers import import_ai_models_config
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = "Import cmbot-style ai_models.json into encrypted AiModel records."
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument("path", help="Path to ai_models.json.")
|
||||
parser.add_argument(
|
||||
"--create-default-aliases",
|
||||
action="store_true",
|
||||
help="Create title-standard and image-standard default aliases when possible.",
|
||||
)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
path = Path(options["path"])
|
||||
if not path.exists():
|
||||
raise CommandError(f"File not found: {path}")
|
||||
|
||||
try:
|
||||
config = json.loads(path.read_text(encoding="utf-8-sig"))
|
||||
result = import_ai_models_config(
|
||||
config,
|
||||
create_default_aliases=options["create_default_aliases"],
|
||||
)
|
||||
except Exception as exc:
|
||||
raise CommandError(str(exc)) from exc
|
||||
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS(
|
||||
"Imported AI models: "
|
||||
f"created={result['created']}, "
|
||||
f"updated={result['updated']}, "
|
||||
f"aliases={result['aliases']}"
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,56 @@
|
||||
# Generated by Django 5.2.15 on 2026-07-02 02:37
|
||||
|
||||
import django.db.models.deletion
|
||||
from django.db import migrations, models
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
initial = True
|
||||
|
||||
dependencies = [
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.CreateModel(
|
||||
name='AiModel',
|
||||
fields=[
|
||||
('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
|
||||
('name', models.CharField(max_length=128, unique=True)),
|
||||
('url', models.URLField(max_length=500)),
|
||||
('model', models.CharField(max_length=128)),
|
||||
('api_type', models.CharField(choices=[('auto', 'Auto'), ('chat', 'Chat completions'), ('gemini', 'Gemini'), ('images', 'Images generations'), ('images_edits', 'Images edits')], default='auto', max_length=32)),
|
||||
('api_key_encrypted', models.TextField(blank=True, editable=False)),
|
||||
('capabilities', models.JSONField(blank=True, default=list)),
|
||||
('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(blank=True, default=dict)),
|
||||
('is_active', models.BooleanField(default=True)),
|
||||
('created_at', models.DateTimeField(auto_now_add=True)),
|
||||
('updated_at', models.DateTimeField(auto_now=True)),
|
||||
],
|
||||
options={
|
||||
'db_table': 'ai_model',
|
||||
'ordering': ('name',),
|
||||
},
|
||||
),
|
||||
migrations.CreateModel(
|
||||
name='ModelAlias',
|
||||
fields=[
|
||||
('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
|
||||
('alias', models.SlugField(max_length=64)),
|
||||
('operation_type', models.CharField(choices=[('title', 'Generate title'), ('image', 'Generate image')], max_length=32)),
|
||||
('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)),
|
||||
('ai_model', models.ForeignKey(on_delete=django.db.models.deletion.PROTECT, related_name='aliases', to='ai.aimodel')),
|
||||
],
|
||||
options={
|
||||
'db_table': 'model_alias',
|
||||
'ordering': ('operation_type', 'alias'),
|
||||
'indexes': [models.Index(fields=['operation_type', 'is_default'], name='model_alias_operati_622f4b_idx'), models.Index(fields=['operation_type', 'is_active'], name='model_alias_operati_32d832_idx')],
|
||||
'constraints': [models.UniqueConstraint(fields=('operation_type', 'alias'), name='unique_model_alias_per_operation')],
|
||||
},
|
||||
),
|
||||
]
|
||||
+133
-1
@@ -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)
|
||||
|
||||
@@ -41,6 +41,12 @@ class ResolvedModel:
|
||||
capabilities = frozenset(
|
||||
item.strip() for item in raw_capabilities.split(",") if item.strip()
|
||||
)
|
||||
elif isinstance(raw_capabilities, Mapping):
|
||||
capabilities = frozenset(
|
||||
str(key)
|
||||
for key, value in raw_capabilities.items()
|
||||
if key in {"text", "image", "vision"} and bool(value)
|
||||
)
|
||||
else:
|
||||
capabilities = frozenset(str(item) for item in raw_capabilities)
|
||||
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from cryptography.fernet import Fernet, InvalidToken
|
||||
from django.conf import settings
|
||||
|
||||
|
||||
class AiKeyEncryptionError(RuntimeError):
|
||||
"""Raised when AI provider key encryption/decryption cannot proceed."""
|
||||
|
||||
|
||||
PREFIX = "fernet:"
|
||||
|
||||
|
||||
def encrypt_api_key(raw_api_key: str) -> str:
|
||||
api_key = str(raw_api_key or "").strip()
|
||||
if not api_key:
|
||||
raise AiKeyEncryptionError("AI api_key cannot be blank")
|
||||
token = _fernet().encrypt(api_key.encode("utf-8")).decode("ascii")
|
||||
return f"{PREFIX}{token}"
|
||||
|
||||
|
||||
def decrypt_api_key(encrypted_api_key: str) -> str:
|
||||
encrypted = str(encrypted_api_key or "")
|
||||
if not encrypted:
|
||||
return ""
|
||||
if not encrypted.startswith(PREFIX):
|
||||
raise AiKeyEncryptionError("AI api_key is not stored in encrypted form")
|
||||
token = encrypted[len(PREFIX) :].encode("ascii")
|
||||
try:
|
||||
return _fernet().decrypt(token).decode("utf-8")
|
||||
except InvalidToken as exc:
|
||||
raise AiKeyEncryptionError("AI api_key cannot be decrypted") from exc
|
||||
|
||||
|
||||
def mask_secret(encrypted_api_key: str) -> str:
|
||||
return "********" if encrypted_api_key else ""
|
||||
|
||||
|
||||
def _fernet() -> Fernet:
|
||||
key = getattr(settings, "AI_KEY_ENCRYPTION_KEY", "")
|
||||
if not key:
|
||||
raise AiKeyEncryptionError("AI_KEY_ENCRYPTION_KEY is not configured")
|
||||
try:
|
||||
return Fernet(str(key).encode("ascii"))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise AiKeyEncryptionError("AI_KEY_ENCRYPTION_KEY is invalid") from exc
|
||||
+251
-1
@@ -1,11 +1,20 @@
|
||||
import base64
|
||||
|
||||
from django.test import SimpleTestCase
|
||||
from cryptography.fernet import Fernet
|
||||
from django.contrib.admin.sites import AdminSite
|
||||
from django.test import RequestFactory, SimpleTestCase, TestCase, override_settings
|
||||
|
||||
from apps.ai.admin import AiModelAdmin
|
||||
from apps.ai.aliases import AliasNotFoundError, ModelCapabilityError, resolve_alias
|
||||
from apps.ai.importers import import_ai_models_config
|
||||
from apps.ai.models import AiModel, ModelAlias
|
||||
from apps.ai.providers import AiCapabilityError, ResolvedModel, get_provider, resolve_api_type
|
||||
from apps.ai.providers.openai_compatible import ChatCompletionsProvider, ImagesEditsProvider
|
||||
|
||||
|
||||
TEST_ENCRYPTION_KEY = Fernet.generate_key().decode("ascii")
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, payload, content=b""):
|
||||
self.payload = payload
|
||||
@@ -175,3 +184,244 @@ class ImagesEditsProviderTests(SimpleTestCase):
|
||||
|
||||
with self.assertRaises(AiCapabilityError):
|
||||
provider.generate_text("Generate title", model)
|
||||
|
||||
|
||||
@override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY)
|
||||
class AiModelEncryptionTests(TestCase):
|
||||
def test_api_key_is_encrypted_and_resolved_model_decrypts_it(self):
|
||||
model = AiModel(
|
||||
name="GPT-5.5 text",
|
||||
url="https://api.vectorengine.ai/v1",
|
||||
model="gpt-5.5",
|
||||
api_type=AiModel.ApiType.CHAT,
|
||||
capabilities=["text"],
|
||||
)
|
||||
model.set_api_key("sk-test-secret")
|
||||
model.save()
|
||||
|
||||
self.assertNotIn("sk-test-secret", model.api_key_encrypted)
|
||||
self.assertTrue(model.api_key_encrypted.startswith("fernet:"))
|
||||
self.assertEqual(model.get_api_key(), "sk-test-secret")
|
||||
|
||||
resolved = model.to_resolved_model()
|
||||
self.assertEqual(resolved.api_key, "sk-test-secret")
|
||||
self.assertEqual(resolved.capabilities, frozenset({"text"}))
|
||||
|
||||
def test_capabilities_accept_dict_shape(self):
|
||||
model = AiModel(
|
||||
name="Vision image",
|
||||
url="https://api.vectorengine.ai/v1/images/edits",
|
||||
model="gpt-image-2",
|
||||
api_type=AiModel.ApiType.IMAGES_EDITS,
|
||||
capabilities={"image": True, "vision": True, "text": False},
|
||||
)
|
||||
|
||||
self.assertEqual(model.capabilities_set(), frozenset({"image", "vision"}))
|
||||
|
||||
|
||||
@override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY)
|
||||
class AliasResolutionTests(TestCase):
|
||||
def setUp(self):
|
||||
self.text_model = self.create_model(
|
||||
name="GPT-5.5 text",
|
||||
model="gpt-5.5",
|
||||
api_type=AiModel.ApiType.CHAT,
|
||||
capabilities=["text", "vision"],
|
||||
)
|
||||
self.image_model = self.create_model(
|
||||
name="GPT Image 2",
|
||||
url="https://api.vectorengine.ai/v1/images/edits",
|
||||
model="gpt-image-2",
|
||||
api_type=AiModel.ApiType.IMAGES_EDITS,
|
||||
capabilities=["image", "vision"],
|
||||
)
|
||||
|
||||
def create_model(
|
||||
self,
|
||||
*,
|
||||
name,
|
||||
model,
|
||||
capabilities,
|
||||
url="https://api.vectorengine.ai/v1",
|
||||
api_type=AiModel.ApiType.CHAT,
|
||||
):
|
||||
ai_model = AiModel(
|
||||
name=name,
|
||||
url=url,
|
||||
model=model,
|
||||
api_type=api_type,
|
||||
capabilities=capabilities,
|
||||
)
|
||||
ai_model.set_api_key("sk-test-secret")
|
||||
ai_model.save()
|
||||
return ai_model
|
||||
|
||||
def test_resolve_default_title_alias(self):
|
||||
ModelAlias.objects.create(
|
||||
operation_type=ModelAlias.OperationType.TITLE,
|
||||
alias="title-standard",
|
||||
ai_model=self.text_model,
|
||||
is_default=True,
|
||||
)
|
||||
|
||||
resolved = resolve_alias(ModelAlias.OperationType.TITLE)
|
||||
|
||||
self.assertEqual(resolved.model, "gpt-5.5")
|
||||
self.assertEqual(resolved.api_key, "sk-test-secret")
|
||||
|
||||
def test_resolve_named_image_alias(self):
|
||||
ModelAlias.objects.create(
|
||||
operation_type=ModelAlias.OperationType.IMAGE,
|
||||
alias="image-edit",
|
||||
ai_model=self.image_model,
|
||||
)
|
||||
|
||||
resolved = resolve_alias(ModelAlias.OperationType.IMAGE, "image-edit")
|
||||
|
||||
self.assertEqual(resolved.model, "gpt-image-2")
|
||||
self.assertIn("image", resolved.capabilities)
|
||||
|
||||
def test_resolve_alias_rejects_capability_mismatch(self):
|
||||
ModelAlias.objects.create(
|
||||
operation_type=ModelAlias.OperationType.TITLE,
|
||||
alias="bad-title",
|
||||
ai_model=self.image_model,
|
||||
)
|
||||
|
||||
with self.assertRaises(ModelCapabilityError):
|
||||
resolve_alias(ModelAlias.OperationType.TITLE, "bad-title")
|
||||
|
||||
def test_resolve_alias_ignores_inactive_aliases(self):
|
||||
ModelAlias.objects.create(
|
||||
operation_type=ModelAlias.OperationType.TITLE,
|
||||
alias="inactive-title",
|
||||
ai_model=self.text_model,
|
||||
is_active=False,
|
||||
)
|
||||
|
||||
with self.assertRaises(AliasNotFoundError):
|
||||
resolve_alias(ModelAlias.OperationType.TITLE, "inactive-title")
|
||||
|
||||
def test_model_alias_save_rejects_second_default_for_operation(self):
|
||||
ModelAlias.objects.create(
|
||||
operation_type=ModelAlias.OperationType.TITLE,
|
||||
alias="title-standard",
|
||||
ai_model=self.text_model,
|
||||
is_default=True,
|
||||
)
|
||||
|
||||
with self.assertRaisesMessage(Exception, "Only one default alias"):
|
||||
ModelAlias.objects.create(
|
||||
operation_type=ModelAlias.OperationType.TITLE,
|
||||
alias="title-backup",
|
||||
ai_model=self.text_model,
|
||||
is_default=True,
|
||||
)
|
||||
|
||||
|
||||
@override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY)
|
||||
class AiModelsImportTests(TestCase):
|
||||
def test_import_cmbot_config_encrypts_keys_and_creates_default_aliases(self):
|
||||
result = import_ai_models_config(
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"name": "Nano Banana 2",
|
||||
"url": "https://api.vectorengine.ai/v1/chat/completions",
|
||||
"model": "gemini-3.1-flash-image-preview",
|
||||
"api_key": "sk-image",
|
||||
"api_type": "auto",
|
||||
"timeout_seconds": 0,
|
||||
"connect_timeout_seconds": 30,
|
||||
"extra_body": {},
|
||||
},
|
||||
{
|
||||
"name": "GPT-5.5 text",
|
||||
"url": "https://api.vectorengine.ai/v1",
|
||||
"model": "gpt-5.5",
|
||||
"api_key": "sk-text",
|
||||
"api_type": "chat",
|
||||
"timeout_seconds": 0,
|
||||
"connect_timeout_seconds": 30,
|
||||
"extra_body": {},
|
||||
},
|
||||
]
|
||||
},
|
||||
create_default_aliases=True,
|
||||
)
|
||||
|
||||
self.assertEqual(result, {"created": 2, "updated": 0, "aliases": 2})
|
||||
text_model = AiModel.objects.get(name="GPT-5.5 text")
|
||||
image_model = AiModel.objects.get(name="Nano Banana 2")
|
||||
self.assertEqual(text_model.capabilities_set(), frozenset({"text", "vision"}))
|
||||
self.assertEqual(image_model.capabilities_set(), frozenset({"image", "vision"}))
|
||||
self.assertNotIn("sk-text", text_model.api_key_encrypted)
|
||||
self.assertEqual(text_model.get_api_key(), "sk-text")
|
||||
|
||||
title_alias = ModelAlias.objects.get(
|
||||
operation_type=ModelAlias.OperationType.TITLE,
|
||||
alias="title-standard",
|
||||
)
|
||||
image_alias = ModelAlias.objects.get(
|
||||
operation_type=ModelAlias.OperationType.IMAGE,
|
||||
alias="image-standard",
|
||||
)
|
||||
self.assertEqual(title_alias.ai_model, text_model)
|
||||
self.assertEqual(image_alias.ai_model, image_model)
|
||||
|
||||
|
||||
@override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY)
|
||||
class AiModelAdminTests(TestCase):
|
||||
def test_admin_form_uses_write_only_api_key_field(self):
|
||||
request = RequestFactory().get("/admin/apps/ai/aimodel/add/")
|
||||
model_admin = AiModelAdmin(AiModel, AdminSite())
|
||||
|
||||
form_class = model_admin.get_form(request)
|
||||
|
||||
self.assertIn("api_key", form_class.base_fields)
|
||||
self.assertNotIn("api_key_encrypted", form_class.base_fields)
|
||||
self.assertFalse(form_class.base_fields["api_key"].widget.render_value)
|
||||
|
||||
def test_admin_form_encrypts_api_key_on_save(self):
|
||||
form_class = AiModelAdmin(AiModel, AdminSite()).form
|
||||
form = form_class(
|
||||
data={
|
||||
"name": "GPT-5.5 text",
|
||||
"url": "https://api.vectorengine.ai/v1",
|
||||
"model": "gpt-5.5",
|
||||
"api_type": AiModel.ApiType.CHAT,
|
||||
"api_key": "sk-admin-secret",
|
||||
"capabilities": '["text"]',
|
||||
"timeout_seconds": 0,
|
||||
"connect_timeout_seconds": 30,
|
||||
"extra_body": "{}",
|
||||
"is_active": "on",
|
||||
}
|
||||
)
|
||||
|
||||
self.assertTrue(form.is_valid(), form.errors)
|
||||
ai_model = form.save()
|
||||
|
||||
self.assertNotIn("sk-admin-secret", ai_model.api_key_encrypted)
|
||||
self.assertEqual(ai_model.get_api_key(), "sk-admin-secret")
|
||||
|
||||
@override_settings(AI_KEY_ENCRYPTION_KEY="")
|
||||
def test_admin_form_reports_missing_encryption_key(self):
|
||||
form_class = AiModelAdmin(AiModel, AdminSite()).form
|
||||
form = form_class(
|
||||
data={
|
||||
"name": "GPT-5.5 text",
|
||||
"url": "https://api.vectorengine.ai/v1",
|
||||
"model": "gpt-5.5",
|
||||
"api_type": AiModel.ApiType.CHAT,
|
||||
"api_key": "sk-admin-secret",
|
||||
"capabilities": '["text"]',
|
||||
"timeout_seconds": 0,
|
||||
"connect_timeout_seconds": 30,
|
||||
"extra_body": "{}",
|
||||
"is_active": "on",
|
||||
}
|
||||
)
|
||||
|
||||
self.assertFalse(form.is_valid())
|
||||
self.assertIn("AI_KEY_ENCRYPTION_KEY is not configured", str(form.errors))
|
||||
|
||||
Reference in New Issue
Block a user