feat(api): generate memoir chapter images asynchronously via Celery

Made-with: Cursor

api/tasks/memoir_tasks.py

@@ -26,9 +26,13 @@ from agents.prompts.memory_prompts import (
     CHAPTER_CATEGORIES,
 )
 from agents.prompts.profile_prompts import format_user_profile_context
+import hashlib
+
 from services.memoir_images.parser import build_initial_image_assets, parse_image_placeholders
 from services.memoir_images.prompting import MemoirImagePromptService
+from services.memoir_images.provider import LiblibImageProvider
 from services.memoir_images.settings import MemoirImageSettings
+from services.memoir_images.storage import TencentCosStorageService
 
 logger = logging.getLogger(__name__)
 
@@ -526,8 +530,71 @@ def generate_chapter_content(self, user_id: str, stage: str, new_content: str):
         raise self.retry(exc=e)
 
 
+def build_cos_key(user_id: str, chapter_id: str, index: int, prompt: str) -> str:
+    short_hash = hashlib.sha1(prompt.encode("utf-8")).hexdigest()[:10]
+    return f"memoirs/{user_id}/{chapter_id}/{index}-{short_hash}.png"
+
+
 @shared_task(bind=True, max_retries=3, default_retry_delay=30)
 def generate_chapter_images(self, chapter_id: str):
     """Async task to generate images for a chapter's pending image assets."""
-    logger.info(f"图片生成任务（桩）: chapter_id={chapter_id}")
-    return {"status": "stub", "chapter_id": chapter_id}
+    db = SessionLocal()
+    try:
+        chapter = db.get(Chapter, chapter_id)
+        if not chapter or not chapter.images:
+            return {"status": "no_images"}
+
+        settings = MemoirImageSettings.from_env()
+        prompt_service = MemoirImagePromptService(llm_service.get_llm(), settings)
+        provider = LiblibImageProvider()
+        storage = TencentCosStorageService.from_env()
+
+        for item in chapter.images:
+            if item.get("status") == "completed" and item.get("url"):
+                continue
+            if item.get("status") not in {"pending", "failed"}:
+                continue
+
+            item["status"] = "processing"
+            db.commit()
+
+            try:
+                context_lines = (chapter.content or "").split("\n")
+                context_excerpt = " ".join(context_lines[:5])[:200]
+
+                prompt_data = prompt_service.build_prompt(
+                    chapter_title=chapter.title,
+                    chapter_category=chapter.category or "",
+                    description=item.get("description", ""),
+                    context_excerpt=context_excerpt,
+                )
+                job = provider.submit_generation(
+                    prompt=prompt_data["prompt"],
+                    size=prompt_data["size"],
+                    style=prompt_data["style"],
+                )
+                if job["status"] != "completed":
+                    job = provider.poll_until_complete(
+                        job,
+                        poll_interval_seconds=settings.poll_interval_seconds,
+                        max_attempts=settings.max_attempts,
+                    )
+                image_bytes = provider.download_image(job)
+                key = build_cos_key(chapter.user_id, chapter.id, item["index"], prompt_data["prompt"])
+                item["url"] = storage.upload_bytes(image_bytes, key, "image/png")
+                item["prompt"] = prompt_data["prompt"]
+                item["style"] = prompt_data["style"]
+                item["size"] = prompt_data["size"]
+                item["status"] = "completed"
+                item["error"] = None
+            except Exception as exc:
+                item["status"] = "failed"
+                item["error"] = str(exc)
+                logger.warning(f"图片生成失败: chapter={chapter_id}, index={item.get('index')}, error={exc}")
+
+            item["updated_at"] = datetime.now(timezone.utc).isoformat()
+
+        db.commit()
+        return {"status": "success"}
+    finally:
+        db.close()

api/tests/test_generate_chapter_images_task.py

@@ -0,0 +1,104 @@
+import unittest
+from unittest.mock import Mock, patch
+
+from api.tasks.memoir_tasks import generate_chapter_images
+
+
+class GenerateChapterImagesTaskTest(unittest.TestCase):
+    @patch("api.tasks.memoir_tasks.SessionLocal")
+    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
+    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
+    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
+    def test_generate_chapter_images_marks_successful_item_completed(
+        self,
+        prompt_service_cls,
+        provider_cls,
+        storage_cls,
+        session_local_cls,
+    ):
+        chapter = type(
+            "ChapterStub",
+            (),
+            {
+                "id": "chapter-1",
+                "user_id": "user-1",
+                "title": "童年的夏天",
+                "category": "childhood",
+                "content": "那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}",
+                "images": [
+                    {
+                        "index": 0,
+                        "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
+                        "description": "南方小镇的青石板路",
+                        "status": "pending",
+                        "url": None,
+                    }
+                ],
+            },
+        )()
+
+        db = Mock()
+        db.get.return_value = chapter
+        session_local_cls.return_value = db
+        prompt_service_cls.return_value.build_prompt.return_value = {
+            "prompt": "A serene southern China town",
+            "style": "watercolor",
+            "size": "1024x1024",
+            "prompt_context": "childhood: 童年的夏天",
+        }
+        provider_inst = provider_cls.return_value
+        provider_inst.submit_generation.return_value = {
+            "status": "completed",
+            "image_url": "https://provider.example.com/1.png",
+        }
+        provider_inst.download_image.return_value = b"png-bytes"
+        storage_inst = storage_cls.from_env.return_value
+        storage_inst.upload_bytes.return_value = "https://cos.example.com/memoirs/u1/c1/0.png"
+
+        generate_chapter_images.run("chapter-1")
+
+        self.assertEqual(chapter.images[0]["status"], "completed")
+        self.assertEqual(chapter.images[0]["url"], "https://cos.example.com/memoirs/u1/c1/0.png")
+        self.assertEqual(chapter.images[0]["prompt"], "A serene southern China town")
+        db.commit.assert_called()
+
+    @patch("api.tasks.memoir_tasks.SessionLocal")
+    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
+    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
+    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
+    def test_generate_chapter_images_skips_completed_items_for_idempotency(
+        self,
+        prompt_service_cls,
+        provider_cls,
+        storage_cls,
+        session_local_cls,
+    ):
+        chapter = type(
+            "ChapterStub",
+            (),
+            {
+                "id": "chapter-1",
+                "user_id": "user-1",
+                "title": "童年的夏天",
+                "category": "childhood",
+                "content": "那条路我一直记得。",
+                "images": [
+                    {
+                        "index": 0,
+                        "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
+                        "description": "南方小镇的青石板路",
+                        "status": "completed",
+                        "url": "https://cos.example.com/already-there.png",
+                    }
+                ],
+            },
+        )()
+
+        db = Mock()
+        db.get.return_value = chapter
+        session_local_cls.return_value = db
+
+        generate_chapter_images.run("chapter-1")
+
+        provider_cls.return_value.submit_generation.assert_not_called()
+        storage_cls.from_env.return_value.upload_bytes.assert_not_called()