feat: 完成T-303 AI批量生成
- 新增 generate_batch，先并发生成标题再并发生成封面，成功逐条 set_generated 落库 - Tab② 接入开始生成、停止、进度展示和双击新旧封面预览 - 新增 GenerateWorker，通过 worker signal 回传进度与行刷新 - 补充批量生成成功、失败、停止取消和 GUI worker 单元测试 - 同步任务看板、API、路由、当前状态与 progress 文档
This commit is contained in:
chengma
+154
-1
@@ -4,13 +4,14 @@ import json
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
|
||||
from _helpers import TempDirMixin
|
||||
|
||||
from app import ai, appconfig
|
||||
from app import ai, appconfig, db
|
||||
|
||||
|
||||
class _Response:
|
||||
@@ -64,6 +65,35 @@ class AITests(TempDirMixin, unittest.TestCase):
|
||||
cfg["ai"]["jpg_quality"] = 80
|
||||
return cfg
|
||||
|
||||
def _collected_tasks(self, temp_dir, cfg, titles=None):
|
||||
titles = titles or ["旧标题A", "旧标题B"]
|
||||
db.init_db(cfg["db_path"])
|
||||
batch_id = db.create_batch([os.path.join(temp_dir, "input.xlsx")], path=cfg["db_path"])
|
||||
db.insert_tasks(
|
||||
batch_id,
|
||||
[
|
||||
{
|
||||
"source_file_abs": os.path.join(temp_dir, "input.xlsx"),
|
||||
"source_sheet": "商品",
|
||||
"source_row": index + 2,
|
||||
"account_name": "Excel主店",
|
||||
"alias": "alias-a",
|
||||
"item_id": "5110063951%s" % index,
|
||||
}
|
||||
for index in range(len(titles))
|
||||
],
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
tasks = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
|
||||
for task, title in zip(tasks, titles):
|
||||
db.set_collected(
|
||||
task.id,
|
||||
title,
|
||||
os.path.join(temp_dir, "%s_old.jpg" % task.item_id),
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
return batch_id, db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
|
||||
|
||||
def test_gen_title_uses_configured_model_and_retries(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
models_path = os.path.join(temp_dir, "ai_models.json")
|
||||
@@ -159,6 +189,129 @@ class AITests(TempDirMixin, unittest.TestCase):
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_generate_batch_persists_titles_and_covers_per_task(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg = self._config()
|
||||
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
|
||||
cfg["image_dir"] = os.path.join(temp_dir, "images")
|
||||
cfg["ai"]["title_concurrency"] = 2
|
||||
cfg["ai"]["image_concurrency"] = 2
|
||||
batch_id, tasks = self._collected_tasks(temp_dir, cfg)
|
||||
cover_prompts = []
|
||||
progress = []
|
||||
|
||||
def fake_title(title_prompt, old_title, **kwargs):
|
||||
self.assertEqual("标题提示", title_prompt)
|
||||
return "新" + old_title
|
||||
|
||||
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
|
||||
cover_prompts.append(cover_prompt)
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
with open(out_path, "wb") as fh:
|
||||
fh.write(b"jpeg")
|
||||
return out_path
|
||||
|
||||
with mock.patch("app.ai.gen_title", side_effect=fake_title), \
|
||||
mock.patch("app.ai.gen_cover", side_effect=fake_cover):
|
||||
summary = ai.generate_batch(
|
||||
tasks,
|
||||
{
|
||||
"title": "标题提示",
|
||||
"cover": "封面 {新标题} {店铺} {商品id}",
|
||||
},
|
||||
ai_cfg={
|
||||
"config": cfg,
|
||||
"db_path": cfg["db_path"],
|
||||
"account_by_alias": {
|
||||
"alias-a": SimpleNamespace(account_name="主店", slug="main")
|
||||
},
|
||||
},
|
||||
on_progress=progress.append,
|
||||
)
|
||||
|
||||
self.assertTrue(summary["ok"])
|
||||
self.assertEqual(2, summary["title_done"])
|
||||
self.assertEqual(2, summary["cover_done"])
|
||||
self.assertEqual(0, summary["failed"])
|
||||
self.assertTrue(progress)
|
||||
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
|
||||
self.assertTrue(all(task.stage == "generated" for task in updated))
|
||||
self.assertEqual({"新旧标题A", "新旧标题B"}, {task.new_title for task in updated})
|
||||
self.assertTrue(all(task.new_cover_path.endswith("_new.jpg") for task in updated))
|
||||
self.assertTrue(all(os.path.exists(task.new_cover_path) for task in updated))
|
||||
self.assertIn("主店", "\n".join(cover_prompts))
|
||||
self.assertIn("新旧标题A", "\n".join(cover_prompts))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_generate_batch_marks_failed_task_without_blocking_others(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg = self._config()
|
||||
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
|
||||
cfg["image_dir"] = os.path.join(temp_dir, "images")
|
||||
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["好标题", "坏标题"])
|
||||
|
||||
def fake_title(title_prompt, old_title, **kwargs):
|
||||
if old_title == "坏标题":
|
||||
raise ai.AIError("标题生成失败")
|
||||
return "新" + old_title
|
||||
|
||||
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
with open(out_path, "wb") as fh:
|
||||
fh.write(b"jpeg")
|
||||
return out_path
|
||||
|
||||
with mock.patch("app.ai.gen_title", side_effect=fake_title), \
|
||||
mock.patch("app.ai.gen_cover", side_effect=fake_cover):
|
||||
summary = ai.generate_batch(
|
||||
tasks,
|
||||
{"title": "标题提示", "cover": "封面 {新标题}"},
|
||||
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
|
||||
)
|
||||
|
||||
self.assertFalse(summary["ok"])
|
||||
self.assertEqual(1, summary["title_done"])
|
||||
self.assertEqual(1, summary["cover_done"])
|
||||
self.assertEqual(1, summary["failed"])
|
||||
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
|
||||
by_title = {task.old_title: task for task in updated}
|
||||
self.assertEqual("generated", by_title["好标题"].stage)
|
||||
self.assertEqual("success", by_title["好标题"].status)
|
||||
self.assertEqual("collected", by_title["坏标题"].stage)
|
||||
self.assertEqual("failed", by_title["坏标题"].status)
|
||||
self.assertIn("标题生成失败", by_title["坏标题"].last_error)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_generate_batch_stop_before_scheduling_keeps_tasks_collected(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg = self._config()
|
||||
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
|
||||
cfg["image_dir"] = os.path.join(temp_dir, "images")
|
||||
batch_id, tasks = self._collected_tasks(temp_dir, cfg)
|
||||
|
||||
with mock.patch("app.ai.gen_title") as gen_title, \
|
||||
mock.patch("app.ai.gen_cover") as gen_cover:
|
||||
summary = ai.generate_batch(
|
||||
tasks,
|
||||
{"title": "标题提示", "cover": "封面"},
|
||||
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
|
||||
should_stop=lambda: True,
|
||||
)
|
||||
|
||||
self.assertFalse(summary["ok"])
|
||||
self.assertTrue(summary["cancelled"])
|
||||
self.assertEqual(0, summary["title_done"])
|
||||
self.assertEqual(0, summary["cover_done"])
|
||||
self.assertEqual(0, summary["failed"])
|
||||
gen_title.assert_not_called()
|
||||
gen_cover.assert_not_called()
|
||||
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
|
||||
self.assertTrue(all(task.stage == "collected" for task in updated))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user