298 lines
10 KiB
Python
298 lines
10 KiB
Python
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
|
|
|
|
from cryptography.fernet import Fernet
|
|
from django.conf import settings
|
|
from django.core.management.base import BaseCommand, CommandError
|
|
from django.db import transaction
|
|
from django.test.utils import override_settings
|
|
|
|
from apps.ai.aliases import resolve_alias
|
|
from apps.ai.models import AiModel, ModelAlias
|
|
from apps.ai.providers import get_provider
|
|
from apps.ai.providers.openai_compatible import ChatCompletionsProvider
|
|
|
|
|
|
DEFAULT_TITLE_PROMPT = "为一件简约白色T恤生成3个中文电商标题"
|
|
RECORDED_TITLE_CONTENT = (
|
|
"1. 简约白T清爽百搭日常通勤短袖\n"
|
|
"2. 纯净白色基础款T恤舒适亲肤上衣\n"
|
|
"3. 夏季百搭白T休闲宽松男女同款"
|
|
)
|
|
RECORDED_IMAGE_BYTES = b"recorded-image-bytes"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SmokeResult:
|
|
mode: str
|
|
operation: str
|
|
alias: str
|
|
model_used: str
|
|
elapsed_ms: int
|
|
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 self.payload
|
|
|
|
def raise_for_status(self) -> None:
|
|
return None
|
|
|
|
|
|
class RecordedSession:
|
|
def __init__(self, payload: dict[str, Any]) -> None:
|
|
self.payload = payload
|
|
self.posts: list[dict[str, Any]] = []
|
|
self.trust_env = True
|
|
|
|
def post(self, url: str, **kwargs: Any) -> RecordedResponse:
|
|
sanitized_kwargs = {
|
|
key: value for key, value in kwargs.items() if key != "headers"
|
|
}
|
|
self.posts.append({"url": url, **sanitized_kwargs})
|
|
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", "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",
|
|
help="Use a recorded provider response and temporary rolled-back config.",
|
|
)
|
|
|
|
def handle(self, *args, **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:
|
|
if not settings.AI_KEY_ENCRYPTION_KEY:
|
|
raise CommandError("AI_KEY_ENCRYPTION_KEY is not configured.")
|
|
return run_title_smoke(
|
|
mode="real",
|
|
alias=options["alias"],
|
|
prompt=options["prompt"],
|
|
resolution=options["resolution"],
|
|
provider=None,
|
|
)
|
|
|
|
def _run_recorded_title(self, options) -> SmokeResult:
|
|
alias = options["alias"] or f"t-104-recorded-title-{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-104 recorded title {uuid4().hex[:8]}",
|
|
url="https://recorded.invalid/v1",
|
|
model="recorded-title-model",
|
|
api_type=AiModel.ApiType.CHAT,
|
|
capabilities=["text"],
|
|
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.TITLE,
|
|
alias=alias,
|
|
ai_model=ai_model,
|
|
)
|
|
result = run_title_smoke(
|
|
mode="recorded",
|
|
alias=alias,
|
|
prompt=options["prompt"],
|
|
resolution=options["resolution"],
|
|
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
|
|
|
|
def _write_result(self, result: SmokeResult) -> None:
|
|
self.stdout.write(self.style.SUCCESS("Smoke generation OK"))
|
|
self.stdout.write(f"mode={result.mode}")
|
|
self.stdout.write(f"operation={result.operation}")
|
|
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}")
|
|
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(
|
|
*,
|
|
mode: str,
|
|
alias: str | None,
|
|
prompt: str,
|
|
resolution: str,
|
|
provider,
|
|
) -> SmokeResult:
|
|
started = perf_counter()
|
|
model = resolve_alias(ModelAlias.OperationType.TITLE, alias)
|
|
selected_provider = provider or get_provider(model.api_type, model.url)
|
|
generation = selected_provider.generate_text(
|
|
prompt,
|
|
model,
|
|
resolution=resolution,
|
|
parameters={"temperature": 0},
|
|
)
|
|
elapsed_ms = int((perf_counter() - started) * 1000)
|
|
return SmokeResult(
|
|
mode=mode,
|
|
operation=ModelAlias.OperationType.TITLE,
|
|
alias=alias or "<default:title>",
|
|
model_used=generation.model_used,
|
|
elapsed_ms=elapsed_ms,
|
|
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
|