fix: address phase 1 ai review
This commit is contained in:
@@ -124,7 +124,7 @@ def _ensure_default_aliases(models: list[AiModel]) -> int:
|
||||
).update(is_default=False)
|
||||
ModelAlias.objects.update_or_create(
|
||||
operation_type=ModelAlias.OperationType.IMAGE,
|
||||
alias="image-standard",
|
||||
alias="image-hd",
|
||||
defaults={"ai_model": image_model, "is_default": True, "is_active": True},
|
||||
)
|
||||
aliases += 1
|
||||
|
||||
@@ -14,7 +14,7 @@ class Command(BaseCommand):
|
||||
parser.add_argument(
|
||||
"--create-default-aliases",
|
||||
action="store_true",
|
||||
help="Create title-standard and image-standard default aliases when possible.",
|
||||
help="Create title-standard and image-hd default aliases when possible.",
|
||||
)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from time import perf_counter
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
@@ -23,6 +25,7 @@ RECORDED_TITLE_CONTENT = (
|
||||
"2. 纯净白色基础款T恤舒适亲肤上衣\n"
|
||||
"3. 夏季百搭白T休闲宽松男女同款"
|
||||
)
|
||||
RECORDED_IMAGE_BYTES = b"recorded-image-bytes"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -32,20 +35,25 @@ class SmokeResult:
|
||||
alias: str
|
||||
model_used: str
|
||||
elapsed_ms: int
|
||||
title_count: int
|
||||
first_title: str
|
||||
title_count: int | None = None
|
||||
first_title: str = ""
|
||||
image_bytes: int | None = None
|
||||
|
||||
|
||||
class RecordedResponse:
|
||||
def __init__(self, payload: dict[str, Any]) -> None:
|
||||
self.payload = payload
|
||||
|
||||
def json(self) -> dict[str, Any]:
|
||||
return {"choices": [{"message": {"content": RECORDED_TITLE_CONTENT}}]}
|
||||
return self.payload
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class RecordedSession:
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, payload: dict[str, Any]) -> None:
|
||||
self.payload = payload
|
||||
self.posts: list[dict[str, Any]] = []
|
||||
self.trust_env = True
|
||||
|
||||
@@ -54,17 +62,19 @@ class RecordedSession:
|
||||
key: value for key, value in kwargs.items() if key != "headers"
|
||||
}
|
||||
self.posts.append({"url": url, **sanitized_kwargs})
|
||||
return RecordedResponse()
|
||||
return RecordedResponse(self.payload)
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = "Run a minimal AI generation smoke through alias resolution and provider code."
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument("operation", choices=("title",))
|
||||
parser.add_argument("operation", choices=("title", "image"))
|
||||
parser.add_argument("--alias", help="Alias to resolve. Defaults to operation default.")
|
||||
parser.add_argument("--prompt", default=DEFAULT_TITLE_PROMPT)
|
||||
parser.add_argument("--resolution", default="1K")
|
||||
parser.add_argument("--image-file", help="Input image path for image/edit smoke.")
|
||||
parser.add_argument("--image-mime-type", default="image/png")
|
||||
parser.add_argument(
|
||||
"--recorded",
|
||||
action="store_true",
|
||||
@@ -72,14 +82,18 @@ class Command(BaseCommand):
|
||||
)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
if options["operation"] != "title":
|
||||
raise CommandError("Only title smoke is implemented.")
|
||||
|
||||
result = (
|
||||
self._run_recorded_title(options)
|
||||
if options["recorded"]
|
||||
else self._run_real_title(options)
|
||||
)
|
||||
if options["operation"] == "title":
|
||||
result = (
|
||||
self._run_recorded_title(options)
|
||||
if options["recorded"]
|
||||
else self._run_real_title(options)
|
||||
)
|
||||
else:
|
||||
result = (
|
||||
self._run_recorded_image(options)
|
||||
if options["recorded"]
|
||||
else self._run_real_image(options)
|
||||
)
|
||||
self._write_result(result)
|
||||
|
||||
def _run_real_title(self, options) -> SmokeResult:
|
||||
@@ -119,7 +133,57 @@ class Command(BaseCommand):
|
||||
alias=alias,
|
||||
prompt=options["prompt"],
|
||||
resolution=options["resolution"],
|
||||
provider=ChatCompletionsProvider(session=RecordedSession()),
|
||||
provider=ChatCompletionsProvider(
|
||||
session=RecordedSession(recorded_title_payload())
|
||||
),
|
||||
)
|
||||
transaction.set_rollback(True)
|
||||
return result
|
||||
|
||||
def _run_real_image(self, options) -> SmokeResult:
|
||||
if not settings.AI_KEY_ENCRYPTION_KEY:
|
||||
raise CommandError("AI_KEY_ENCRYPTION_KEY is not configured.")
|
||||
return run_image_smoke(
|
||||
mode="real",
|
||||
alias=options["alias"],
|
||||
prompt=options["prompt"],
|
||||
resolution=options["resolution"],
|
||||
image=read_image_file(options["image_file"]),
|
||||
image_mime_type=options["image_mime_type"],
|
||||
provider=None,
|
||||
)
|
||||
|
||||
def _run_recorded_image(self, options) -> SmokeResult:
|
||||
alias = options["alias"] or f"t-105-recorded-image-{uuid4().hex[:8]}"
|
||||
encryption_key = Fernet.generate_key().decode("ascii")
|
||||
with override_settings(AI_KEY_ENCRYPTION_KEY=encryption_key):
|
||||
with transaction.atomic():
|
||||
ai_model = AiModel(
|
||||
name=f"T-105 recorded image {uuid4().hex[:8]}",
|
||||
url="https://recorded.invalid/v1/chat/completions",
|
||||
model="recorded-image-model",
|
||||
api_type=AiModel.ApiType.CHAT,
|
||||
capabilities=["image", "vision"],
|
||||
timeout_seconds=0,
|
||||
connect_timeout_seconds=30,
|
||||
)
|
||||
ai_model.set_api_key("sk-recorded-placeholder")
|
||||
ai_model.save()
|
||||
ModelAlias.objects.create(
|
||||
operation_type=ModelAlias.OperationType.IMAGE,
|
||||
alias=alias,
|
||||
ai_model=ai_model,
|
||||
)
|
||||
result = run_image_smoke(
|
||||
mode="recorded",
|
||||
alias=alias,
|
||||
prompt=options["prompt"],
|
||||
resolution=options["resolution"],
|
||||
image=b"recorded-input-image",
|
||||
image_mime_type=options["image_mime_type"],
|
||||
provider=ChatCompletionsProvider(
|
||||
session=RecordedSession(recorded_image_payload())
|
||||
),
|
||||
)
|
||||
transaction.set_rollback(True)
|
||||
return result
|
||||
@@ -131,8 +195,11 @@ class Command(BaseCommand):
|
||||
self.stdout.write(f"alias={result.alias}")
|
||||
self.stdout.write(f"model_used={result.model_used}")
|
||||
self.stdout.write(f"elapsed_ms={result.elapsed_ms}")
|
||||
self.stdout.write(f"title_count={result.title_count}")
|
||||
self.stdout.write(f"first_title={result.first_title}")
|
||||
if result.title_count is not None:
|
||||
self.stdout.write(f"title_count={result.title_count}")
|
||||
self.stdout.write(f"first_title={result.first_title}")
|
||||
if result.image_bytes is not None:
|
||||
self.stdout.write(f"image_bytes={result.image_bytes}")
|
||||
|
||||
|
||||
def run_title_smoke(
|
||||
@@ -162,3 +229,69 @@ def run_title_smoke(
|
||||
title_count=len(generation.titles),
|
||||
first_title=generation.titles[0] if generation.titles else generation.text,
|
||||
)
|
||||
|
||||
|
||||
def run_image_smoke(
|
||||
*,
|
||||
mode: str,
|
||||
alias: str | None,
|
||||
prompt: str,
|
||||
resolution: str,
|
||||
image: bytes | None,
|
||||
image_mime_type: str,
|
||||
provider,
|
||||
) -> SmokeResult:
|
||||
started = perf_counter()
|
||||
model = resolve_alias(ModelAlias.OperationType.IMAGE, alias)
|
||||
selected_provider = provider or get_provider(model.api_type, model.url)
|
||||
generation = selected_provider.generate_image(
|
||||
prompt,
|
||||
model,
|
||||
image=image,
|
||||
image_mime_type=image_mime_type,
|
||||
resolution=resolution,
|
||||
parameters={"temperature": 0},
|
||||
)
|
||||
elapsed_ms = int((perf_counter() - started) * 1000)
|
||||
return SmokeResult(
|
||||
mode=mode,
|
||||
operation=ModelAlias.OperationType.IMAGE,
|
||||
alias=alias or "<default:image>",
|
||||
model_used=generation.model_used,
|
||||
elapsed_ms=elapsed_ms,
|
||||
image_bytes=len(generation.image),
|
||||
)
|
||||
|
||||
|
||||
def recorded_title_payload() -> dict[str, Any]:
|
||||
return {"choices": [{"message": {"content": RECORDED_TITLE_CONTENT}}]}
|
||||
|
||||
|
||||
def recorded_image_payload() -> dict[str, Any]:
|
||||
encoded = base64.b64encode(RECORDED_IMAGE_BYTES).decode("ascii")
|
||||
return {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": [
|
||||
{"type": "text", "text": "done"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{encoded}",
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def read_image_file(path: str | None) -> bytes | None:
|
||||
if not path:
|
||||
return None
|
||||
try:
|
||||
return Path(path).read_bytes()
|
||||
except OSError as exc:
|
||||
raise CommandError(f"Unable to read image file: {path}") from exc
|
||||
|
||||
@@ -28,6 +28,35 @@ from .utils import (
|
||||
)
|
||||
|
||||
|
||||
SAFE_PARAMETER_KEYS = frozenset(
|
||||
{
|
||||
"temperature",
|
||||
"top_p",
|
||||
"topP",
|
||||
"top_k",
|
||||
"topK",
|
||||
"max_tokens",
|
||||
"max_output_tokens",
|
||||
"maxOutputTokens",
|
||||
"presence_penalty",
|
||||
"presencePenalty",
|
||||
"frequency_penalty",
|
||||
"frequencyPenalty",
|
||||
"seed",
|
||||
"stop",
|
||||
}
|
||||
)
|
||||
SAFE_GENERATION_CONFIG_KEYS = frozenset(
|
||||
{
|
||||
"temperature",
|
||||
"topP",
|
||||
"topK",
|
||||
"maxOutputTokens",
|
||||
"stopSequences",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class BaseHttpProvider:
|
||||
def __init__(self, session: requests.Session | None = None):
|
||||
self.session = session or requests.Session()
|
||||
@@ -371,6 +400,20 @@ def apply_extra_body(
|
||||
model: ResolvedModel,
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> None:
|
||||
payload.update(model.extra_body)
|
||||
apply_safe_parameters(payload, model.extra_body)
|
||||
if parameters:
|
||||
payload.update(dict(parameters))
|
||||
apply_safe_parameters(payload, parameters)
|
||||
|
||||
|
||||
def apply_safe_parameters(payload: dict[str, Any], values: Mapping[str, Any]) -> None:
|
||||
for key, value in values.items():
|
||||
if key == "generationConfig" and isinstance(value, Mapping):
|
||||
generation_config = payload.setdefault("generationConfig", {})
|
||||
if not isinstance(generation_config, dict):
|
||||
continue
|
||||
for config_key, config_value in value.items():
|
||||
if config_key in SAFE_GENERATION_CONFIG_KEYS:
|
||||
generation_config[config_key] = config_value
|
||||
continue
|
||||
if key in SAFE_PARAMETER_KEYS:
|
||||
payload[key] = value
|
||||
|
||||
@@ -110,14 +110,15 @@ def decode_image_data_url(data_url: str) -> bytes:
|
||||
|
||||
|
||||
def resolution_to_size(resolution: str) -> str:
|
||||
key = str(resolution).strip().upper()
|
||||
mapping = {
|
||||
"512": "512x512",
|
||||
"512px": "512x512",
|
||||
"512PX": "512x512",
|
||||
"1K": "1024x1024",
|
||||
"2K": "2048x2048",
|
||||
"4K": "4096x4096",
|
||||
}
|
||||
return mapping.get(str(resolution), str(resolution))
|
||||
return mapping.get(key, str(resolution).strip())
|
||||
|
||||
|
||||
def extract_image_from_response(
|
||||
|
||||
+119
-2
@@ -13,6 +13,7 @@ from apps.ai.importers import import_ai_models_config
|
||||
from apps.ai.models import AiConfigAuditLog, AiModel, ModelAlias
|
||||
from apps.ai.providers import AiCapabilityError, ResolvedModel, get_provider, resolve_api_type
|
||||
from apps.ai.providers.openai_compatible import ChatCompletionsProvider, ImagesEditsProvider
|
||||
from apps.ai.providers.utils import resolution_to_size
|
||||
|
||||
|
||||
TEST_ENCRYPTION_KEY = Fernet.generate_key().decode("ascii")
|
||||
@@ -54,6 +55,12 @@ class ProviderRegistryTests(SimpleTestCase):
|
||||
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")
|
||||
|
||||
|
||||
class ChatCompletionsProviderTests(SimpleTestCase):
|
||||
def test_generate_text_builds_chat_payload_and_cleans_titles(self):
|
||||
session = FakeSession(
|
||||
@@ -92,6 +99,54 @@ class ChatCompletionsProviderTests(SimpleTestCase):
|
||||
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")
|
||||
@@ -172,6 +227,44 @@ class ImagesEditsProviderTests(SimpleTestCase):
|
||||
("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(
|
||||
@@ -367,7 +460,7 @@ class AiModelsImportTests(TestCase):
|
||||
)
|
||||
image_alias = ModelAlias.objects.get(
|
||||
operation_type=ModelAlias.OperationType.IMAGE,
|
||||
alias="image-standard",
|
||||
alias="image-hd",
|
||||
)
|
||||
self.assertEqual(title_alias.ai_model, text_model)
|
||||
self.assertEqual(image_alias.ai_model, image_model)
|
||||
@@ -523,7 +616,7 @@ class AiConfigAuditAdminTests(TestCase):
|
||||
)
|
||||
alias = ModelAlias.objects.create(
|
||||
operation_type=ModelAlias.OperationType.IMAGE,
|
||||
alias="image-standard",
|
||||
alias="image-hd",
|
||||
ai_model=text_model,
|
||||
)
|
||||
alias.ai_model = image_model
|
||||
@@ -588,3 +681,27 @@ class AiGenerationSmokeCommandTests(TestCase):
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user