fix: address phase 1 ai review

2026-07-02 14:47:22 +08:00
parent 2d9c233683
commit 2f93374193
18 changed files with 499 additions and 56 deletions
@@ -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