feat: add ai model alias configuration
This commit is contained in:
QiuSW
+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