feat(suite): confirm AI writing cost before request

This commit is contained in:
chengma
2026-07-17 11:48:25 +08:00
parent a06a02ddcf
commit 1ef5721a20
11 changed files with 620 additions and 18 deletions
+69
View File
@@ -8,6 +8,9 @@ from app import cmhub_models
class CMHubModelDisplayTests(unittest.TestCase):
def tearDown(self):
cmhub_models.clear_model_catalog_cache()
def test_alias_label_prefers_display_name_tier_and_points(self):
label = cmhub_models.alias_label(
{
@@ -44,6 +47,72 @@ class CMHubModelDisplayTests(unittest.TestCase):
self.assertIn("生图别名 image-hd", summary)
self.assertIn("扣点以返回结果为准", summary)
def test_catalog_price_only_accepts_one_unconditional_matching_vision_model(self):
models = [
{
"alias": "vision-standard",
"operation_type": "vision",
"requires_image": True,
"pricing_status": "priced",
"prices": [{"points_cost": 2}],
}
]
cmhub_models.cache_model_catalog("https://CMHUB.example.com/", models, now=100)
cached = cmhub_models.cached_model_catalog(
"https://cmhub.example.com",
"vision-standard",
now=120,
)
self.assertEqual("2", cmhub_models.format_points_cost(
cmhub_models.unambiguous_points_cost(
cached,
"vision-standard",
"vision",
requires_image=True,
)
))
def test_catalog_price_rejects_conditional_missing_or_ambiguous_prices(self):
base = {
"alias": "vision-standard",
"operation_type": "vision",
"requires_image": True,
"pricing_status": "priced",
}
cases = [
dict(base, prices=[{"resolution": "1K", "points_cost": 2}]),
dict(base, prices=[{"points_cost": 2}, {"points_cost": 3}]),
dict(base, pricing_status="unpriced", prices=[{"points_cost": 2}]),
dict(base, requires_image=False, prices=[{"points_cost": 2}]),
dict(base, prices=[{}]),
]
for model in cases:
with self.subTest(model=model):
self.assertIsNone(
cmhub_models.unambiguous_points_cost(
[model],
"vision-standard",
"vision",
requires_image=True,
)
)
cmhub_models.cache_model_catalog(
"https://cmhub.example.com",
[dict(base, prices=[{"points_cost": 2}])],
now=100,
)
self.assertIsNone(
cmhub_models.cached_model_catalog(
"https://cmhub.example.com",
"vision-standard",
max_age_seconds=20,
now=121,
)
)
if __name__ == "__main__":
unittest.main()
+132 -3
View File
@@ -9,7 +9,15 @@ sys.path.insert(0, os.path.dirname(__file__))
from _helpers import TempDirMixin
from app import accounts, appconfig, image_studio, image_studio_images, product_suite, prompts
from app import (
accounts,
appconfig,
cmhub_models,
image_studio,
image_studio_images,
product_suite,
prompts,
)
from app import gui
if gui.QT_IMPORT_ERROR is not None:
@@ -965,7 +973,14 @@ class ProductSuiteGuiTests(TempDirMixin, unittest.TestCase):
_AiWriteWorker,
), mock.patch.object(tab, "_start_thread", return_value=object()), mock.patch.object(
tab, "_status"
) as status:
) as status, mock.patch.object(
tab,
"_confirm_ai_write_request",
side_effect=lambda target, asset_ids, points_cost: tab._start_confirmed_ai_write(
target,
asset_ids,
),
):
tab.start_ai_write()
self.assertEqual(
@@ -985,6 +1000,113 @@ class ProductSuiteGuiTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_ai_write_uses_cached_price_before_confirming_the_first_eight_images(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
config["ai"] = {
"cmhub": {
"base_url": "https://cmhub.example.com",
"vision_alias": "vision-standard",
}
}
appconfig.save_cmhub_config({"api_key": "test-key"}, path=config["cmhub_config_path"])
project, assets = self._create_project_with_assets(temp_dir, config, 9)
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
self.addCleanup(tab.close)
self.addCleanup(cmhub_models.clear_model_catalog_cache)
state = tab._displayed_state
state.account_alias = "alias-a"
state.item_id = project.item_id
state.project_id = project.id
state.project_binding_state = project.binding_state
tab._load_state(state)
cmhub_models.cache_model_catalog(
"https://cmhub.example.com",
[
{
"alias": "vision-standard",
"operation_type": "vision",
"requires_image": True,
"pricing_status": "priced",
"prices": [{"points_cost": 2}],
}
],
)
with mock.patch.object(tab, "_confirm_ai_write_request") as confirm:
tab.start_ai_write()
state_arg, asset_ids, points_cost = confirm.call_args.args
self.assertIs(state_arg, state)
self.assertEqual(tuple(asset.id for asset in assets[:8]), asset_ids)
self.assertEqual("2", cmhub_models.format_points_cost(points_cost))
self.assertIsNone(state.ai_worker)
self.assertIsNone(state.ai_price_worker)
self.assert_removed(temp_dir)
def test_ai_write_cancelled_confirmation_does_not_start_worker(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
project, assets = self._create_project_with_assets(temp_dir, config, 1)
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
self.addCleanup(tab.close)
state = tab._displayed_state
state.account_alias = "alias-a"
state.item_id = project.item_id
state.project_id = project.id
state.project_binding_state = project.binding_state
tab._load_state(state)
class _Button:
pass
class _MessageBox:
Question = 1
AcceptRole = 2
RejectRole = 3
def __init__(self, *args):
self.start_button = _Button()
self.cancel_button = _Button()
self.clicked = self.cancel_button
def setIcon(self, value):
pass
def setWindowTitle(self, value):
pass
def setText(self, value):
pass
def addButton(self, text, role):
return self.start_button if role == self.AcceptRole else self.cancel_button
def setDefaultButton(self, button):
pass
def setEscapeButton(self, button):
pass
def exec(self):
return 0
def clickedButton(self):
return self.clicked
with mock.patch("app.gui.tabs.product_suite.QMessageBox", _MessageBox), mock.patch.object(
tab,
"_start_confirmed_ai_write",
) as start_confirmed:
tab._confirm_ai_write_request(state, (assets[0].id,), None)
start_confirmed.assert_not_called()
self.assertIsNone(state.ai_worker)
self.assertFalse(state.ai_confirmation_open)
self.assert_removed(temp_dir)
def test_prompt_edit_expands_shrinks_and_reflows_without_internal_scrollbars(self):
edit = AutoHeightPlainTextEdit()
self.addCleanup(edit.close)
@@ -2458,7 +2580,14 @@ class ProductSuiteGuiTests(TempDirMixin, unittest.TestCase):
with mock.patch(
"app.gui.tabs.product_suite.ProductSuiteAiWriteWorker",
_AiWriteWorker,
), mock.patch.object(tab, "_start_thread", return_value=object()):
), mock.patch.object(tab, "_start_thread", return_value=object()), mock.patch.object(
tab,
"_confirm_ai_write_request",
side_effect=lambda target, asset_ids, points_cost: tab._start_confirmed_ai_write(
target,
asset_ids,
),
):
tab.start_ai_write()
self.assertIn("未绑定商品", captured["context"])
self.assertNotIn("draft_", captured["context"])
+31 -1
View File
@@ -8,7 +8,7 @@ os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
from _helpers import REPO_ROOT # noqa: F401
from app import db, image_studio, image_studio_images, workers
from app import cmhub_models, db, image_studio, image_studio_images, workers
if workers.QT_IMPORT_ERROR is not None:
raise unittest.SkipTest("PySide6 未安装")
@@ -18,6 +18,7 @@ from PySide6.QtWidgets import QApplication
from app.workers import BaseWorker, run_worker
from app.gui.workers import (
CMHubModelCatalogWorker,
ImageStudioDownloadOriginalWorker,
ImageStudioPullImagesWorker,
ProductSuiteAiWriteWorker,
@@ -230,6 +231,35 @@ class WorkerTests(unittest.TestCase):
gen_title.assert_not_called()
self.assertEqual(expected, result)
def test_cmhub_model_catalog_worker_caches_models_without_exposing_key(self):
cmhub_models.clear_model_catalog_cache()
self.addCleanup(cmhub_models.clear_model_catalog_cache)
worker = CMHubModelCatalogWorker(
"https://cmhub.example.com/",
"sk-cmhub-secret",
connect_timeout=7,
use_system_proxy=True,
)
models = [{"alias": "vision-standard", "operation_type": "vision"}]
with mock.patch("app.gui.workers.ai.fetch_cmhub_models", return_value=models) as fetch:
result = worker.execute()
fetch.assert_called_once_with(
"https://cmhub.example.com",
"sk-cmhub-secret",
connect_timeout=7,
use_system_proxy=True,
)
self.assertEqual(models, result["models"])
self.assertEqual(
models,
cmhub_models.cached_model_catalog(
"https://cmhub.example.com",
"vision-standard",
),
)
def test_product_suite_history_export_worker_uses_round_export_service(self):
worker = ProductSuiteHistoryExportWorker(
7,