135 lines
4.4 KiB
Python
135 lines
4.4 KiB
Python
"""Prompt file management and variable rendering."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
from . import appconfig
|
|
|
|
TITLE_PROMPT_PATH = appconfig.title_prompt_path()
|
|
COVER_PROMPTS_DIR = appconfig.cover_prompts_dir()
|
|
TEMPLATE_EXT = ".txt"
|
|
INVALID_NAME_CHARS = set('\\/:*?"<>|')
|
|
|
|
|
|
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 list_cover_templates(directory=COVER_PROMPTS_DIR):
|
|
"""Return cover 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_cover_template(name, directory=COVER_PROMPTS_DIR) -> str:
|
|
"""Load one cover 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_cover_template(name, text, directory=COVER_PROMPTS_DIR) -> None:
|
|
"""Save one cover 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_cover_template(old, new, directory=COVER_PROMPTS_DIR) -> None:
|
|
"""Rename a cover 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_cover_template(name, directory=COVER_PROMPTS_DIR) -> None:
|
|
"""Delete one cover prompt template."""
|
|
|
|
path = _template_path(name, directory)
|
|
if not os.path.exists(path):
|
|
raise PromptError(f"封面提示词模板不存在: {_normalize_name(name)}")
|
|
os.remove(path)
|
|
|
|
|
|
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
|