890 lines
31 KiB
Python
890 lines
31 KiB
Python
import base64
|
|
import io
|
|
|
|
from cryptography.fernet import Fernet
|
|
from django.contrib.admin.sites import AdminSite
|
|
from django.contrib.auth import get_user_model
|
|
from django.core.management import call_command
|
|
from django.test import RequestFactory, SimpleTestCase, TestCase, override_settings
|
|
|
|
from apps.ai.admin import AiConfigAuditLogAdmin, AiModelAdmin, ModelAliasAdmin
|
|
from apps.ai.aliases import AliasNotFoundError, ModelCapabilityError, resolve_alias
|
|
from apps.ai.importers import import_ai_models_config
|
|
from apps.ai.models import AiConfigAuditLog, AiModel, ModelAlias
|
|
from apps.ai.providers import (
|
|
AiCapabilityError,
|
|
MultimodalImage,
|
|
ResolvedModel,
|
|
get_provider,
|
|
resolve_api_type,
|
|
)
|
|
from apps.ai.providers.openai_compatible import (
|
|
ChatCompletionsProvider,
|
|
GeminiProvider,
|
|
ImagesEditsProvider,
|
|
)
|
|
from apps.ai.providers.utils import image_request_timeout, resolution_to_size
|
|
|
|
|
|
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 ProviderUtilsTests(SimpleTestCase):
|
|
def test_resolution_to_size_normalizes_case(self):
|
|
self.assertEqual(resolution_to_size("1k"), "1024x1024")
|
|
self.assertEqual(resolution_to_size("512px"), "512x512")
|
|
|
|
@override_settings(AI_IMAGE_UPSTREAM_DEADLINE_SECONDS=180)
|
|
def test_image_request_timeout_caps_read_timeout_to_deadline(self):
|
|
self.assertEqual(image_request_timeout(30, 0, "4K"), (30, 180))
|
|
self.assertEqual(image_request_timeout(30, 120, "4K"), (30, 120))
|
|
|
|
|
|
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_parameters_cannot_override_core_chat_payload_fields(self):
|
|
session = FakeSession(
|
|
FakeResponse(
|
|
{
|
|
"choices": [
|
|
{"message": {"content": "1. Safe Title"}}
|
|
]
|
|
}
|
|
)
|
|
)
|
|
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",
|
|
extra_body={
|
|
"model": "bad-extra-model",
|
|
"stream": True,
|
|
"temperature": 0.1,
|
|
},
|
|
)
|
|
|
|
provider.generate_text(
|
|
"Generate titles",
|
|
model,
|
|
parameters={
|
|
"model": "gpt-image-2",
|
|
"n": 5,
|
|
"messages": [],
|
|
"stream": True,
|
|
"size": "4096x4096",
|
|
"temperature": 0.2,
|
|
},
|
|
)
|
|
|
|
payload = session.posts[0]["json"]
|
|
self.assertEqual(payload["model"], "gpt-5.5")
|
|
self.assertFalse(payload["stream"])
|
|
self.assertEqual(
|
|
payload["messages"],
|
|
[{"role": "user", "content": [{"type": "text", "text": "Generate titles"}]}],
|
|
)
|
|
self.assertNotIn("n", payload)
|
|
self.assertNotIn("size", payload)
|
|
self.assertEqual(payload["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,"))
|
|
|
|
@override_settings(AI_IMAGE_UPSTREAM_DEADLINE_SECONDS=180)
|
|
def test_generate_image_caps_post_and_download_timeouts(self):
|
|
session = FakeSession(
|
|
FakeResponse(
|
|
{
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": [
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "https://cdn.example.com/out.png"},
|
|
}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
),
|
|
FakeResponse({}, content=b"generated-image"),
|
|
)
|
|
provider = ChatCompletionsProvider(session=session)
|
|
model = ResolvedModel(
|
|
name="Slow image",
|
|
url="https://api.vectorengine.ai/v1/chat/completions",
|
|
model="image-model",
|
|
api_key="test-key",
|
|
api_type="chat",
|
|
timeout_seconds=0,
|
|
connect_timeout_seconds=30,
|
|
)
|
|
|
|
result = provider.generate_image("Generate product image", model, resolution="4K")
|
|
|
|
self.assertEqual(result.image, b"generated-image")
|
|
self.assertEqual(session.posts[0]["timeout"], (30, 180))
|
|
self.assertEqual(session.gets[0]["timeout"], 180)
|
|
|
|
@override_settings(AI_IMAGE_UPSTREAM_DEADLINE_SECONDS=180)
|
|
def test_generate_text_does_not_use_image_deadline(self):
|
|
session = FakeSession(
|
|
FakeResponse({"choices": [{"message": {"content": "1. Red Dress"}}]})
|
|
)
|
|
provider = ChatCompletionsProvider(session=session)
|
|
model = ResolvedModel(
|
|
name="Slow text",
|
|
url="https://api.vectorengine.ai/v1/chat/completions",
|
|
model="text-model",
|
|
api_key="test-key",
|
|
api_type="chat",
|
|
timeout_seconds=0,
|
|
connect_timeout_seconds=30,
|
|
)
|
|
|
|
provider.generate_text("Generate titles", model, resolution="4K")
|
|
|
|
self.assertEqual(session.posts[0]["timeout"], (30, 600))
|
|
|
|
def test_analyze_images_builds_ordered_payload_and_preserves_full_text(self):
|
|
session = FakeSession(
|
|
FakeResponse(
|
|
{
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": "第一张是正面图。\n第二张是细节图。"
|
|
}
|
|
}
|
|
]
|
|
}
|
|
)
|
|
)
|
|
provider = ChatCompletionsProvider(session=session)
|
|
model = ResolvedModel(
|
|
name="Vision text",
|
|
url="https://api.vectorengine.ai/v1",
|
|
model="vision-model",
|
|
api_key="test-key",
|
|
api_type="chat",
|
|
)
|
|
|
|
result = provider.analyze_images(
|
|
"比较两张商品图",
|
|
model,
|
|
images=(
|
|
MultimodalImage(b"first-image", "image/jpeg"),
|
|
MultimodalImage(b"second-image", "image/png"),
|
|
),
|
|
parameters={"temperature": 0.2, "messages": []},
|
|
)
|
|
|
|
self.assertEqual(result.text, "第一张是正面图。\n第二张是细节图。")
|
|
self.assertEqual(result.titles, ())
|
|
content = session.posts[0]["json"]["messages"][0]["content"]
|
|
self.assertEqual(content[0], {"type": "text", "text": "比较两张商品图"})
|
|
self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/jpeg;base64,"))
|
|
self.assertTrue(content[2]["image_url"]["url"].startswith("data:image/png;base64,"))
|
|
self.assertEqual(session.posts[0]["json"]["temperature"], 0.2)
|
|
|
|
def test_gemini_analyze_images_builds_ordered_inline_data(self):
|
|
session = FakeSession(
|
|
FakeResponse(
|
|
{
|
|
"candidates": [
|
|
{"content": {"parts": [{"text": "多图分析结果"}]}}
|
|
]
|
|
}
|
|
)
|
|
)
|
|
provider = GeminiProvider(session=session)
|
|
model = ResolvedModel(
|
|
name="Gemini vision",
|
|
url="https://gemini.example.com",
|
|
model="gemini-vision",
|
|
api_key="test-key",
|
|
api_type="gemini",
|
|
)
|
|
|
|
result = provider.analyze_images(
|
|
"理解这些图片",
|
|
model,
|
|
images=(
|
|
MultimodalImage(b"one", "image/webp"),
|
|
MultimodalImage(b"two", "image/jpeg"),
|
|
),
|
|
)
|
|
|
|
self.assertEqual(result.text, "多图分析结果")
|
|
parts = session.posts[0]["json"]["contents"][0]["parts"]
|
|
self.assertEqual(parts[0], {"text": "理解这些图片"})
|
|
self.assertEqual(parts[1]["inlineData"]["mimeType"], "image/webp")
|
|
self.assertEqual(parts[2]["inlineData"]["mimeType"], "image/jpeg")
|
|
self.assertEqual(
|
|
session.posts[0]["json"]["generationConfig"]["responseModalities"],
|
|
["TEXT"],
|
|
)
|
|
|
|
|
|
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_parameters_cannot_override_core_images_edits_fields(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",
|
|
extra_body={
|
|
"model": "bad-extra-model",
|
|
"n": "9",
|
|
"size": "4096x4096",
|
|
"temperature": 0.3,
|
|
},
|
|
)
|
|
|
|
provider.generate_image(
|
|
"Replace background",
|
|
model,
|
|
image=b"source-image",
|
|
resolution="1k",
|
|
parameters={
|
|
"model": "bad-parameter-model",
|
|
"n": "3",
|
|
"size": "1x1",
|
|
"temperature": 0.7,
|
|
},
|
|
)
|
|
|
|
data = session.posts[0]["data"]
|
|
self.assertEqual(data["model"], "gpt-image-2")
|
|
self.assertEqual(data["n"], "1")
|
|
self.assertEqual(data["size"], "1024x1024")
|
|
self.assertEqual(data["temperature"], 0.7)
|
|
|
|
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)
|
|
|
|
with self.assertRaises(AiCapabilityError):
|
|
provider.analyze_images(
|
|
"Analyze image",
|
|
model,
|
|
images=(MultimodalImage(b"source-image"),),
|
|
)
|
|
|
|
|
|
@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_vision_alias_requires_text_and_vision_capabilities(self):
|
|
ModelAlias.objects.create(
|
|
operation_type=ModelAlias.OperationType.VISION,
|
|
alias="vision-standard",
|
|
ai_model=self.text_model,
|
|
is_default=True,
|
|
)
|
|
|
|
resolved = resolve_alias(ModelAlias.OperationType.VISION)
|
|
|
|
self.assertEqual(resolved.model, "gpt-5.5")
|
|
self.assertTrue({"text", "vision"}.issubset(resolved.capabilities))
|
|
|
|
def test_resolve_vision_alias_rejects_vision_model_without_text(self):
|
|
ModelAlias.objects.create(
|
|
operation_type=ModelAlias.OperationType.VISION,
|
|
alias="bad-vision",
|
|
ai_model=self.image_model,
|
|
)
|
|
|
|
with self.assertRaises(ModelCapabilityError):
|
|
resolve_alias(ModelAlias.OperationType.VISION, "bad-vision")
|
|
|
|
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-hd",
|
|
)
|
|
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))
|
|
|
|
|
|
@override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY)
|
|
class AiConfigAuditAdminTests(TestCase):
|
|
def setUp(self):
|
|
self.site = AdminSite()
|
|
self.request = RequestFactory().post("/admin/")
|
|
self.request.user = get_user_model().objects.create_superuser(
|
|
username="auditor",
|
|
email="auditor@example.com",
|
|
password="password",
|
|
)
|
|
|
|
def create_ai_model(
|
|
self,
|
|
*,
|
|
name="GPT-5.5 text",
|
|
model="gpt-5.5",
|
|
capabilities=None,
|
|
api_type=AiModel.ApiType.CHAT,
|
|
url="https://api.vectorengine.ai/v1",
|
|
api_key="sk-test-secret",
|
|
):
|
|
ai_model = AiModel(
|
|
name=name,
|
|
url=url,
|
|
model=model,
|
|
api_type=api_type,
|
|
capabilities=capabilities or ["text"],
|
|
)
|
|
ai_model.set_api_key(api_key)
|
|
ai_model.save()
|
|
return ai_model
|
|
|
|
def test_aimodel_admin_create_writes_sanitized_audit_log(self):
|
|
ai_model = AiModel(
|
|
name="GPT-5.5 text",
|
|
url="https://api.vectorengine.ai/v1",
|
|
model="gpt-5.5",
|
|
api_type=AiModel.ApiType.CHAT,
|
|
capabilities=["text"],
|
|
)
|
|
ai_model.set_api_key("sk-created-secret")
|
|
AiModelAdmin(AiModel, self.site).save_model(
|
|
self.request,
|
|
ai_model,
|
|
form=None,
|
|
change=False,
|
|
)
|
|
|
|
log = AiConfigAuditLog.objects.get()
|
|
self.assertEqual(log.actor, self.request.user)
|
|
self.assertEqual(log.action, AiConfigAuditLog.Action.CREATE)
|
|
self.assertEqual(log.target_type, AiConfigAuditLog.TargetType.AI_MODEL)
|
|
self.assertIn("api_key", log.changed_fields)
|
|
self.assertEqual(log.changes["api_key"], {"old": "empty", "new": "set"})
|
|
self.assertNotIn("sk-created-secret", str(log.changes))
|
|
self.assertNotIn("fernet:", str(log.changes))
|
|
|
|
def test_aimodel_admin_update_logs_field_and_key_changes(self):
|
|
ai_model = self.create_ai_model()
|
|
ai_model.url = "https://api.vectorengine.ai/v2"
|
|
ai_model.set_api_key("sk-new-secret")
|
|
|
|
AiModelAdmin(AiModel, self.site).save_model(
|
|
self.request,
|
|
ai_model,
|
|
form=None,
|
|
change=True,
|
|
)
|
|
|
|
log = AiConfigAuditLog.objects.get()
|
|
self.assertEqual(log.action, AiConfigAuditLog.Action.UPDATE)
|
|
self.assertEqual(set(log.changed_fields), {"url", "api_key"})
|
|
self.assertEqual(
|
|
log.changes["url"],
|
|
{
|
|
"old": "https://api.vectorengine.ai/v1",
|
|
"new": "https://api.vectorengine.ai/v2",
|
|
},
|
|
)
|
|
self.assertEqual(log.changes["api_key"], {"old": "set", "new": "set"})
|
|
self.assertNotIn("sk-new-secret", str(log.changes))
|
|
|
|
def test_model_alias_admin_update_logs_mapping_change(self):
|
|
text_model = self.create_ai_model(name="Text model", model="gpt-5.5")
|
|
image_model = self.create_ai_model(
|
|
name="Image model",
|
|
model="gpt-image-2",
|
|
capabilities=["image"],
|
|
api_type=AiModel.ApiType.IMAGES_EDITS,
|
|
url="https://api.vectorengine.ai/v1/images/edits",
|
|
)
|
|
alias = ModelAlias.objects.create(
|
|
operation_type=ModelAlias.OperationType.IMAGE,
|
|
alias="image-hd",
|
|
ai_model=text_model,
|
|
)
|
|
alias.ai_model = image_model
|
|
|
|
ModelAliasAdmin(ModelAlias, self.site).save_model(
|
|
self.request,
|
|
alias,
|
|
form=None,
|
|
change=True,
|
|
)
|
|
|
|
log = AiConfigAuditLog.objects.get()
|
|
self.assertEqual(log.target_type, AiConfigAuditLog.TargetType.MODEL_ALIAS)
|
|
self.assertEqual(log.changed_fields, ["ai_model"])
|
|
self.assertEqual(
|
|
log.changes["ai_model"],
|
|
{"old": text_model.id, "new": image_model.id},
|
|
)
|
|
|
|
def test_aimodel_admin_delete_writes_audit_log(self):
|
|
ai_model = self.create_ai_model()
|
|
target_id = ai_model.id
|
|
|
|
AiModelAdmin(AiModel, self.site).delete_model(self.request, ai_model)
|
|
|
|
log = AiConfigAuditLog.objects.get()
|
|
self.assertEqual(log.action, AiConfigAuditLog.Action.DELETE)
|
|
self.assertEqual(log.target_id, target_id)
|
|
self.assertIn("api_key", log.changed_fields)
|
|
self.assertEqual(log.changes["api_key"], {"old": "set", "new": "empty"})
|
|
|
|
def test_audit_log_admin_is_read_only(self):
|
|
model_admin = AiConfigAuditLogAdmin(AiConfigAuditLog, self.site)
|
|
|
|
self.assertFalse(model_admin.has_add_permission(self.request))
|
|
self.assertFalse(model_admin.has_change_permission(self.request))
|
|
self.assertFalse(model_admin.has_delete_permission(self.request))
|
|
self.assertIn("changes", model_admin.get_readonly_fields(self.request))
|
|
|
|
|
|
class AiGenerationSmokeCommandTests(TestCase):
|
|
def test_recorded_title_smoke_runs_through_alias_and_provider_without_persisting(self):
|
|
output = io.StringIO()
|
|
|
|
call_command(
|
|
"smoke_ai_generation",
|
|
"title",
|
|
"--recorded",
|
|
"--prompt",
|
|
"为测试商品生成标题",
|
|
stdout=output,
|
|
)
|
|
|
|
text = output.getvalue()
|
|
self.assertIn("Smoke generation OK", text)
|
|
self.assertIn("mode=recorded", text)
|
|
self.assertIn("operation=title", text)
|
|
self.assertIn("alias=t-104-recorded-title-", text)
|
|
self.assertIn("model_used=recorded-title-model", text)
|
|
self.assertIn("title_count=3", text)
|
|
self.assertNotIn("sk-recorded-placeholder", text)
|
|
self.assertNotIn("Bearer", text)
|
|
self.assertEqual(AiModel.objects.count(), 0)
|
|
self.assertEqual(ModelAlias.objects.count(), 0)
|
|
|
|
def test_recorded_image_smoke_runs_through_alias_and_provider_without_persisting(self):
|
|
output = io.StringIO()
|
|
|
|
call_command(
|
|
"smoke_ai_generation",
|
|
"image",
|
|
"--recorded",
|
|
"--prompt",
|
|
"生成测试图片",
|
|
stdout=output,
|
|
)
|
|
|
|
text = output.getvalue()
|
|
self.assertIn("Smoke generation OK", text)
|
|
self.assertIn("mode=recorded", text)
|
|
self.assertIn("operation=image", text)
|
|
self.assertIn("alias=t-105-recorded-image-", text)
|
|
self.assertIn("model_used=recorded-image-model", text)
|
|
self.assertIn("image_bytes=20", text)
|
|
self.assertNotIn("sk-recorded-placeholder", text)
|
|
self.assertNotIn("Bearer", text)
|
|
self.assertEqual(AiModel.objects.count(), 0)
|
|
self.assertEqual(ModelAlias.objects.count(), 0)
|