428 lines
15 KiB
Python
428 lines
15 KiB
Python
import base64
|
|
|
|
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
|
|
self.content = content
|
|
|
|
def json(self):
|
|
return self.payload
|
|
|
|
def raise_for_status(self):
|
|
return None
|
|
|
|
|
|
class FakeSession:
|
|
def __init__(self, *responses):
|
|
self.responses = list(responses)
|
|
self.posts = []
|
|
self.gets = []
|
|
self.trust_env = True
|
|
|
|
def post(self, url, **kwargs):
|
|
self.posts.append({"url": url, **kwargs})
|
|
return self.responses.pop(0)
|
|
|
|
def get(self, url, **kwargs):
|
|
self.gets.append({"url": url, **kwargs})
|
|
return self.responses.pop(0)
|
|
|
|
|
|
class ProviderRegistryTests(SimpleTestCase):
|
|
def test_auto_api_type_resolves_chat_provider_from_url(self):
|
|
url = "https://api.vectorengine.ai/v1/chat/completions"
|
|
|
|
self.assertEqual(resolve_api_type("auto", url), "chat")
|
|
self.assertIsInstance(get_provider("auto", url), ChatCompletionsProvider)
|
|
|
|
|
|
class ChatCompletionsProviderTests(SimpleTestCase):
|
|
def test_generate_text_builds_chat_payload_and_cleans_titles(self):
|
|
session = FakeSession(
|
|
FakeResponse(
|
|
{
|
|
"choices": [
|
|
{"message": {"content": "1. Red Dress\n2. Blue Coat"}}
|
|
]
|
|
}
|
|
)
|
|
)
|
|
provider = ChatCompletionsProvider(session=session)
|
|
model = ResolvedModel(
|
|
name="GPT-5.5 text",
|
|
url="https://api.vectorengine.ai/v1",
|
|
model="gpt-5.5",
|
|
api_key="test-key",
|
|
api_type="chat",
|
|
)
|
|
|
|
result = provider.generate_text(
|
|
"Generate titles",
|
|
model,
|
|
parameters={"temperature": 0.2},
|
|
)
|
|
|
|
self.assertEqual(result.text, "Red Dress")
|
|
self.assertEqual(result.titles, ("Red Dress", "Blue Coat"))
|
|
request = session.posts[0]
|
|
self.assertEqual(
|
|
request["url"],
|
|
"https://api.vectorengine.ai/v1/chat/completions",
|
|
)
|
|
self.assertEqual(request["headers"]["Authorization"], "Bearer test-key")
|
|
self.assertEqual(request["json"]["model"], "gpt-5.5")
|
|
self.assertFalse(request["json"]["stream"])
|
|
self.assertEqual(request["json"]["temperature"], 0.2)
|
|
|
|
def test_generate_image_parses_chat_multimodal_data_url(self):
|
|
generated = b"generated-image"
|
|
encoded = base64.b64encode(generated).decode("ascii")
|
|
session = FakeSession(
|
|
FakeResponse(
|
|
{
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": [
|
|
{"type": "text", "text": "done"},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{encoded}"
|
|
},
|
|
},
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
)
|
|
)
|
|
provider = ChatCompletionsProvider(session=session)
|
|
model = ResolvedModel(
|
|
name="Nano Banana 2",
|
|
url="https://api.vectorengine.ai/v1/chat/completions",
|
|
model="gemini-3.1-flash-image-preview",
|
|
api_key="test-key",
|
|
api_type="auto",
|
|
)
|
|
|
|
result = provider.generate_image(
|
|
"Generate product image",
|
|
model,
|
|
image=b"input-image",
|
|
image_mime_type="image/jpeg",
|
|
)
|
|
|
|
self.assertEqual(result.image, generated)
|
|
content = session.posts[0]["json"]["messages"][0]["content"]
|
|
self.assertEqual(content[0], {"type": "text", "text": "Generate product image"})
|
|
self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/jpeg;base64,"))
|
|
|
|
|
|
class ImagesEditsProviderTests(SimpleTestCase):
|
|
def test_generate_image_builds_multipart_request_and_parses_base64(self):
|
|
generated = b"edited-image"
|
|
encoded = base64.b64encode(generated).decode("ascii")
|
|
session = FakeSession(FakeResponse({"data": [{"b64_json": encoded}]}))
|
|
provider = ImagesEditsProvider(session=session)
|
|
model = ResolvedModel(
|
|
name="GPT Image 2",
|
|
url="https://api.vectorengine.ai/v1/images/edits",
|
|
model="gpt-image-2",
|
|
api_key="test-key",
|
|
api_type="images_edits",
|
|
)
|
|
|
|
result = provider.generate_image(
|
|
"Replace background",
|
|
model,
|
|
image=b"source-image",
|
|
image_mime_type="image/png",
|
|
image_filename="source.png",
|
|
resolution="1K",
|
|
)
|
|
|
|
self.assertEqual(result.image, generated)
|
|
request = session.posts[0]
|
|
self.assertEqual(request["url"], "https://api.vectorengine.ai/v1/images/edits")
|
|
self.assertEqual(request["headers"]["Authorization"], "Bearer test-key")
|
|
self.assertEqual(request["data"]["model"], "gpt-image-2")
|
|
self.assertEqual(request["data"]["size"], "1024x1024")
|
|
self.assertEqual(
|
|
request["files"]["image"],
|
|
("source.png", b"source-image", "image/png"),
|
|
)
|
|
|
|
def test_images_edits_requires_input_image(self):
|
|
provider = ImagesEditsProvider(session=FakeSession())
|
|
model = ResolvedModel(
|
|
name="GPT Image 2",
|
|
url="https://api.vectorengine.ai/v1/images/edits",
|
|
model="gpt-image-2",
|
|
api_key="test-key",
|
|
api_type="images_edits",
|
|
)
|
|
|
|
with self.assertRaises(AiCapabilityError):
|
|
provider.generate_image("Replace background", model)
|
|
|
|
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))
|