413 lines
13 KiB
Python
413 lines
13 KiB
Python
"""Prompt file management and variable rendering."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import tempfile
|
||
from importlib import resources
|
||
|
||
from . import appconfig
|
||
|
||
TITLE_PROMPT_PATH = appconfig.title_prompt_path()
|
||
TITLE_TEMPLATES_DIR = appconfig.title_templates_dir()
|
||
COVER_PROMPTS_DIR = appconfig.cover_prompts_dir()
|
||
IMAGE_STUDIO_PROMPTS_DIR = appconfig.image_studio_prompts_dir()
|
||
PRODUCT_SUITE_PROMPT_PATH = appconfig.product_suite_prompt_path()
|
||
TEMPLATE_EXT = ".txt"
|
||
INVALID_NAME_CHARS = set('\\/:*?"<>|')
|
||
DEFAULT_PROMPTS_PACKAGE = "app.default_prompts"
|
||
DEFAULT_COVER_PROMPTS_PACKAGE = "app.default_prompts.cover"
|
||
DEFAULT_PRODUCT_SUITE_PROMPTS_PACKAGE = "app.default_prompts.product_suite"
|
||
DEFAULT_PRODUCT_SUITE_PROMPT_NAME = "base.txt"
|
||
PRODUCT_SUITE_DEFAULT_ERROR = "内置套图提示词模板无效，请重新安装软件或联系技术支持。"
|
||
|
||
|
||
class PromptError(RuntimeError):
|
||
"""Raised when prompt files or template names are invalid."""
|
||
|
||
|
||
def load_title_prompt(path=TITLE_PROMPT_PATH) -> str:
|
||
"""Load the title prompt text. Missing file means an empty prompt."""
|
||
|
||
if not os.path.exists(path):
|
||
return ""
|
||
with open(path, "r", encoding="utf-8") as fh:
|
||
return fh.read()
|
||
|
||
|
||
def save_title_prompt(text, path=TITLE_PROMPT_PATH) -> None:
|
||
"""Save the title prompt text as UTF-8."""
|
||
|
||
directory = os.path.dirname(os.path.abspath(path))
|
||
if directory:
|
||
os.makedirs(directory, exist_ok=True)
|
||
with open(path, "w", encoding="utf-8") as fh:
|
||
fh.write(str(text or ""))
|
||
|
||
|
||
def ensure_default_prompts(
|
||
title_prompt_path=TITLE_PROMPT_PATH,
|
||
cover_prompts_dir=COVER_PROMPTS_DIR,
|
||
title_templates_dir=None,
|
||
) -> None:
|
||
"""Seed bundled default prompts into an empty user data directory.
|
||
|
||
Existing user prompts are never overwritten. Defaults live in the packaged
|
||
app and are copied into data/ only when the corresponding user files are
|
||
missing, or the title prompt file is blank.
|
||
"""
|
||
|
||
if not _has_non_empty_file(title_prompt_path):
|
||
default_title = _read_default_prompt_text("title_prompt.txt")
|
||
if default_title:
|
||
save_title_prompt(default_title, title_prompt_path)
|
||
|
||
if title_templates_dir is None:
|
||
title_templates_dir = _title_templates_dir_for_prompt_path(title_prompt_path)
|
||
if not list_templates(title_templates_dir):
|
||
default_title = _read_default_prompt_text("title_prompt.txt")
|
||
if default_title:
|
||
save_template("默认", default_title, title_templates_dir)
|
||
|
||
if not list_cover_templates(cover_prompts_dir):
|
||
for name, text in _iter_default_cover_templates():
|
||
save_cover_template(name, text, cover_prompts_dir)
|
||
|
||
|
||
def read_product_suite_prompt(path=PRODUCT_SUITE_PROMPT_PATH) -> str:
|
||
"""Read a product-suite template without hiding invalid user content."""
|
||
|
||
if not os.path.exists(path):
|
||
return ""
|
||
try:
|
||
with open(path, "r", encoding="utf-8") as fh:
|
||
return fh.read()
|
||
except (OSError, UnicodeError) as exc:
|
||
raise PromptError("套图提示词模板读取失败：%s" % exc) from exc
|
||
|
||
|
||
def load_product_suite_prompt(path=PRODUCT_SUITE_PROMPT_PATH) -> str:
|
||
"""Load and validate the current user product-suite template."""
|
||
|
||
text = read_product_suite_prompt(path)
|
||
if not text.strip():
|
||
raise PromptError("套图提示词模板不存在或为空")
|
||
_validate_product_suite_prompt(text, "套图提示词模板无效")
|
||
return text
|
||
|
||
|
||
def load_default_product_suite_prompt() -> str:
|
||
"""Load and validate the packaged product-suite template."""
|
||
|
||
try:
|
||
text = (
|
||
resources.files(DEFAULT_PRODUCT_SUITE_PROMPTS_PACKAGE)
|
||
.joinpath(DEFAULT_PRODUCT_SUITE_PROMPT_NAME)
|
||
.read_text(encoding="utf-8")
|
||
)
|
||
except (FileNotFoundError, ModuleNotFoundError, OSError, UnicodeError) as exc:
|
||
raise PromptError(PRODUCT_SUITE_DEFAULT_ERROR) from exc
|
||
try:
|
||
_validate_product_suite_prompt(text, "内置套图提示词模板无效")
|
||
except PromptError as exc:
|
||
raise PromptError(PRODUCT_SUITE_DEFAULT_ERROR) from exc
|
||
return text
|
||
|
||
|
||
def save_product_suite_prompt(text, path=PRODUCT_SUITE_PROMPT_PATH) -> None:
|
||
"""Validate and atomically save the user product-suite template."""
|
||
|
||
_validate_product_suite_prompt(text, "套图提示词模板无效")
|
||
_atomic_write_text(path, str(text))
|
||
|
||
|
||
def ensure_default_product_suite_prompt(path=PRODUCT_SUITE_PROMPT_PATH) -> str:
|
||
"""Seed a validated packaged template without overwriting user content."""
|
||
|
||
if os.path.exists(path):
|
||
current = read_product_suite_prompt(path)
|
||
if current.strip():
|
||
_validate_product_suite_prompt(current, "套图提示词模板无效")
|
||
return current
|
||
default_text = load_default_product_suite_prompt()
|
||
save_product_suite_prompt(default_text, path)
|
||
return default_text
|
||
|
||
|
||
def restore_default_product_suite_prompt(path=PRODUCT_SUITE_PROMPT_PATH) -> str:
|
||
"""Validate and atomically restore the packaged product-suite template."""
|
||
|
||
default_text = load_default_product_suite_prompt()
|
||
save_product_suite_prompt(default_text, path)
|
||
return default_text
|
||
|
||
|
||
def list_templates(directory):
|
||
"""Return prompt template names sorted by display name."""
|
||
|
||
if not os.path.isdir(directory):
|
||
return []
|
||
names = []
|
||
for filename in os.listdir(directory):
|
||
if filename.lower().endswith(TEMPLATE_EXT):
|
||
names.append(filename[: -len(TEMPLATE_EXT)])
|
||
return sorted(names, key=str.casefold)
|
||
|
||
|
||
def load_template(name, directory) -> str:
|
||
"""Load one prompt template."""
|
||
|
||
path = _template_path(name, directory)
|
||
if not os.path.exists(path):
|
||
raise PromptError(f"提示词模板不存在: {_normalize_name(name)}")
|
||
with open(path, "r", encoding="utf-8") as fh:
|
||
return fh.read()
|
||
|
||
|
||
def save_template(name, text, directory) -> None:
|
||
"""Save one prompt template as UTF-8."""
|
||
|
||
path = _template_path(name, directory)
|
||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||
with open(path, "w", encoding="utf-8") as fh:
|
||
fh.write(str(text or ""))
|
||
|
||
|
||
def rename_template(old, new, directory) -> None:
|
||
"""Rename a prompt template with duplicate-name protection."""
|
||
|
||
old_path = _template_path(old, directory)
|
||
new_path = _template_path(new, directory)
|
||
if not os.path.exists(old_path):
|
||
raise PromptError(f"提示词模板不存在: {_normalize_name(old)}")
|
||
if os.path.exists(new_path):
|
||
raise PromptError(f"提示词模板已存在: {_normalize_name(new)}")
|
||
os.makedirs(os.path.dirname(new_path), exist_ok=True)
|
||
os.replace(old_path, new_path)
|
||
|
||
|
||
def delete_template(name, directory) -> None:
|
||
"""Delete one prompt template."""
|
||
|
||
path = _template_path(name, directory)
|
||
if not os.path.exists(path):
|
||
raise PromptError(f"提示词模板不存在: {_normalize_name(name)}")
|
||
os.remove(path)
|
||
|
||
|
||
def list_title_templates(directory=TITLE_TEMPLATES_DIR):
|
||
"""Return title prompt template names sorted by display name."""
|
||
|
||
return list_templates(directory)
|
||
|
||
|
||
def load_title_template(name, directory=TITLE_TEMPLATES_DIR) -> str:
|
||
"""Load one title prompt template."""
|
||
|
||
return load_template(name, directory)
|
||
|
||
|
||
def save_title_template(name, text, directory=TITLE_TEMPLATES_DIR) -> None:
|
||
"""Save one title prompt template as UTF-8."""
|
||
|
||
save_template(name, text, directory)
|
||
|
||
|
||
def rename_title_template(old, new, directory=TITLE_TEMPLATES_DIR) -> None:
|
||
"""Rename a title prompt template with duplicate-name protection."""
|
||
|
||
rename_template(old, new, directory)
|
||
|
||
|
||
def delete_title_template(name, directory=TITLE_TEMPLATES_DIR) -> None:
|
||
"""Delete one title prompt template."""
|
||
|
||
delete_template(name, directory)
|
||
|
||
|
||
def list_image_studio_templates(directory=IMAGE_STUDIO_PROMPTS_DIR):
|
||
"""Return AI studio prompt template names sorted by display name."""
|
||
|
||
return list_templates(directory)
|
||
|
||
|
||
def load_image_studio_template(name, directory=IMAGE_STUDIO_PROMPTS_DIR) -> str:
|
||
"""Load one AI studio prompt template as raw text."""
|
||
|
||
return load_template(name, directory)
|
||
|
||
|
||
def save_image_studio_template(name, text, directory=IMAGE_STUDIO_PROMPTS_DIR) -> None:
|
||
"""Save one AI studio prompt template as raw UTF-8 text."""
|
||
|
||
save_template(name, text, directory)
|
||
|
||
|
||
def rename_image_studio_template(old, new, directory=IMAGE_STUDIO_PROMPTS_DIR) -> None:
|
||
"""Rename an AI studio prompt template with duplicate-name protection."""
|
||
|
||
rename_template(old, new, directory)
|
||
|
||
|
||
def delete_image_studio_template(name, directory=IMAGE_STUDIO_PROMPTS_DIR) -> None:
|
||
"""Delete one AI studio prompt template."""
|
||
|
||
delete_template(name, directory)
|
||
|
||
|
||
def list_cover_templates(directory=COVER_PROMPTS_DIR):
|
||
"""Return cover template names sorted by display name."""
|
||
|
||
return list_templates(directory)
|
||
|
||
|
||
def load_cover_template(name, directory=COVER_PROMPTS_DIR) -> str:
|
||
"""Load one cover prompt template."""
|
||
|
||
return load_template(name, directory)
|
||
|
||
|
||
def save_cover_template(name, text, directory=COVER_PROMPTS_DIR) -> None:
|
||
"""Save one cover prompt template as UTF-8."""
|
||
|
||
save_template(name, text, directory)
|
||
|
||
|
||
def rename_cover_template(old, new, directory=COVER_PROMPTS_DIR) -> None:
|
||
"""Rename a cover prompt template with duplicate-name protection."""
|
||
|
||
rename_template(old, new, directory)
|
||
|
||
|
||
def delete_cover_template(name, directory=COVER_PROMPTS_DIR) -> None:
|
||
"""Delete one cover prompt template."""
|
||
|
||
delete_template(name, directory)
|
||
|
||
|
||
def render_prompt(template_text, task) -> str:
|
||
"""Render known task variables in a prompt template."""
|
||
|
||
values = {
|
||
"旧标题": _task_value(task, "old_title", "旧标题"),
|
||
"新标题": _task_value(task, "new_title", "新标题"),
|
||
"商品id": _task_value(task, "item_id", "商品id", "商品ID"),
|
||
"店铺": _task_value(task, "shop", "account_name", "店铺", "alias"),
|
||
}
|
||
rendered = str(template_text or "")
|
||
for name, value in values.items():
|
||
rendered = rendered.replace("{" + name + "}", value)
|
||
return rendered
|
||
|
||
|
||
def _task_value(task, *names) -> str:
|
||
for name in names:
|
||
value = None
|
||
if isinstance(task, dict):
|
||
value = task.get(name)
|
||
else:
|
||
value = getattr(task, name, None)
|
||
if value is not None:
|
||
return str(value)
|
||
return ""
|
||
|
||
|
||
def _template_path(name, directory) -> str:
|
||
normalized = _normalize_name(name)
|
||
return os.path.abspath(os.path.join(directory, normalized + TEMPLATE_EXT))
|
||
|
||
|
||
def _normalize_name(name) -> str:
|
||
value = str(name or "").strip()
|
||
if value.lower().endswith(TEMPLATE_EXT):
|
||
value = value[: -len(TEMPLATE_EXT)]
|
||
value = value.strip()
|
||
if not value:
|
||
raise PromptError("提示词模板名不能为空")
|
||
if value in {".", ".."} or any(char in INVALID_NAME_CHARS for char in value):
|
||
raise PromptError(f"提示词模板名非法: {value}")
|
||
if os.path.basename(value) != value:
|
||
raise PromptError(f"提示词模板名非法: {value}")
|
||
return value
|
||
|
||
|
||
def _title_templates_dir_for_prompt_path(title_prompt_path) -> str:
|
||
data_root = os.path.dirname(os.path.abspath(title_prompt_path))
|
||
return os.path.join(data_root, "prompts", "title")
|
||
|
||
|
||
def _has_non_empty_file(path) -> bool:
|
||
try:
|
||
if not os.path.exists(path):
|
||
return False
|
||
with open(path, "r", encoding="utf-8") as fh:
|
||
return bool(fh.read().strip())
|
||
except OSError:
|
||
return False
|
||
|
||
|
||
def _validate_product_suite_prompt(text, prefix) -> None:
|
||
from . import product_suite
|
||
|
||
try:
|
||
product_suite.validate_product_suite_prompt(text)
|
||
except product_suite.ProductSuitePromptError as exc:
|
||
raise PromptError("%s：%s" % (prefix, exc)) from exc
|
||
|
||
|
||
def _atomic_write_text(path, text) -> None:
|
||
target = os.path.abspath(path)
|
||
directory = os.path.dirname(target)
|
||
if directory:
|
||
os.makedirs(directory, exist_ok=True)
|
||
temporary_path = ""
|
||
try:
|
||
with tempfile.NamedTemporaryFile(
|
||
mode="w",
|
||
encoding="utf-8",
|
||
dir=directory or None,
|
||
prefix=".prompt-",
|
||
suffix=".tmp",
|
||
delete=False,
|
||
) as handle:
|
||
temporary_path = handle.name
|
||
handle.write(str(text))
|
||
handle.flush()
|
||
os.fsync(handle.fileno())
|
||
os.replace(temporary_path, target)
|
||
except OSError as exc:
|
||
raise PromptError("套图提示词模板保存失败：%s" % exc) from exc
|
||
finally:
|
||
if temporary_path and os.path.exists(temporary_path):
|
||
try:
|
||
os.remove(temporary_path)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _read_default_prompt_text(filename) -> str:
|
||
try:
|
||
return (
|
||
resources.files(DEFAULT_PROMPTS_PACKAGE)
|
||
.joinpath(filename)
|
||
.read_text(encoding="utf-8")
|
||
)
|
||
except (FileNotFoundError, ModuleNotFoundError, OSError):
|
||
return ""
|
||
|
||
|
||
def _iter_default_cover_templates():
|
||
try:
|
||
root = resources.files(DEFAULT_COVER_PROMPTS_PACKAGE)
|
||
except (ModuleNotFoundError, OSError):
|
||
return
|
||
for entry in root.iterdir():
|
||
if not entry.is_file() or not entry.name.lower().endswith(TEMPLATE_EXT):
|
||
continue
|
||
name = entry.name[: -len(TEMPLATE_EXT)]
|
||
try:
|
||
text = entry.read_text(encoding="utf-8")
|
||
except OSError:
|
||
continue
|
||
yield name, text
|