From 879466fde16f8a3c37205b97af21c5d459a958ba Mon Sep 17 00:00:00 2001 From: Kevin Date: Tue, 10 Mar 2026 16:05:13 +0800 Subject: [PATCH] feat(api): generate memoir chapter images asynchronously via Celery Made-with: Cursor --- api/tasks/memoir_tasks.py | 71 +++++++++++- .../test_generate_chapter_images_task.py | 104 ++++++++++++++++++ 2 files changed, 173 insertions(+), 2 deletions(-) create mode 100644 api/tests/test_generate_chapter_images_task.py diff --git a/api/tasks/memoir_tasks.py b/api/tasks/memoir_tasks.py index 8831da1..3e5421e 100644 --- a/api/tasks/memoir_tasks.py +++ b/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() diff --git a/api/tests/test_generate_chapter_images_task.py b/api/tests/test_generate_chapter_images_task.py new file mode 100644 index 0000000..08cc1d0 --- /dev/null +++ b/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()