cmhub/apps/ai/management/commands/smoke_ai_generation.py

from __future__ import annotations

from dataclasses import dataclass
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休闲宽松男女同款"
)


@dataclass(frozen=True)
class SmokeResult:
    mode: str
    operation: str
    alias: str
    model_used: str
    elapsed_ms: int
    title_count: int
    first_title: str


class RecordedResponse:
    def json(self) -> dict[str, Any]:
        return {"choices": [{"message": {"content": RECORDED_TITLE_CONTENT}}]}

    def raise_for_status(self) -> None:
        return None


class RecordedSession:
    def __init__(self) -> None:
        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()


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("--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(
            "--recorded",
            action="store_true",
            help="Use a recorded provider response and temporary rolled-back config.",
        )

    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)
        )
        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()),
                )
                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}")
        self.stdout.write(f"title_count={result.title_count}")
        self.stdout.write(f"first_title={result.first_title}")


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,
    )