331 lines
11 KiB
Python
331 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""记录 MAOMOMO 文章配图任务状态。
|
||
|
||
本脚本用于把图片生成进度写入状态文件，避免只用聊天记录声称已完成。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
from datetime import datetime, timezone
|
||
import json
|
||
from pathlib import Path
|
||
import shutil
|
||
import sys
|
||
from typing import Any, Dict, Iterable, List, Optional, Tuple
|
||
|
||
|
||
JOBS_FILE = "slide_jobs.json"
|
||
STATE_FILE = "slide_run_state.json"
|
||
|
||
|
||
def now_iso() -> str:
|
||
return datetime.now(timezone.utc).isoformat()
|
||
|
||
|
||
def die(message: str, code: int = 1) -> None:
|
||
print(f"错误：{message}", file=sys.stderr)
|
||
raise SystemExit(code)
|
||
|
||
|
||
def read_text(path: Path) -> str:
|
||
return path.read_text(encoding="utf-8")
|
||
|
||
|
||
def write_json(path: Path, data: Any) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
with path.open("w", encoding="utf-8", newline="\n") as handle:
|
||
json.dump(data, handle, ensure_ascii=False, indent=2)
|
||
handle.write("\n")
|
||
|
||
|
||
def read_json(path: Path) -> Any:
|
||
try:
|
||
return json.loads(read_text(path))
|
||
except FileNotFoundError:
|
||
die(f"找不到文件：{path}")
|
||
except json.JSONDecodeError as exc:
|
||
die(f"JSON 格式错误：{path}: {exc}")
|
||
|
||
|
||
def project_path(project: str) -> Path:
|
||
path = Path(project).expanduser().resolve()
|
||
if not path.exists():
|
||
die(f"项目目录不存在：{path}")
|
||
if not path.is_dir():
|
||
die(f"不是项目目录：{path}")
|
||
return path
|
||
|
||
|
||
def resolve_in_project(project: Path, value: str) -> Path:
|
||
path = Path(value)
|
||
if not path.is_absolute():
|
||
path = project / path
|
||
return path.resolve()
|
||
|
||
|
||
def load_job_doc(project: Path) -> Dict[str, Any]:
|
||
raw = read_json(project / JOBS_FILE)
|
||
if isinstance(raw, list):
|
||
return {"version": 1, "jobs": raw}
|
||
if not isinstance(raw, dict):
|
||
die(f"{JOBS_FILE} 必须是对象或数组。")
|
||
jobs = raw.get("jobs")
|
||
if not isinstance(jobs, list):
|
||
die(f"{JOBS_FILE} 缺少 jobs 数组。")
|
||
return raw
|
||
|
||
|
||
def load_state_doc(project: Path) -> Dict[str, Any]:
|
||
path = project / STATE_FILE
|
||
if not path.exists():
|
||
return {"version": 1, "events": []}
|
||
raw = read_json(path)
|
||
if not isinstance(raw, dict):
|
||
die(f"{STATE_FILE} 必须是对象。")
|
||
if not isinstance(raw.get("events"), list):
|
||
raw["events"] = []
|
||
return raw
|
||
|
||
|
||
def save_docs(project: Path, job_doc: Dict[str, Any], state_doc: Dict[str, Any]) -> None:
|
||
write_json(project / JOBS_FILE, job_doc)
|
||
write_json(project / STATE_FILE, state_doc)
|
||
|
||
|
||
def job_key(job: Dict[str, Any]) -> Optional[str]:
|
||
for key in ("id", "job_id", "slide", "name"):
|
||
value = job.get(key)
|
||
if isinstance(value, str) and value:
|
||
return value
|
||
return None
|
||
|
||
|
||
def find_job(job_doc: Dict[str, Any], wanted: str) -> Dict[str, Any]:
|
||
for job in job_doc["jobs"]:
|
||
if isinstance(job, dict) and job_key(job) == wanted:
|
||
return job
|
||
die(f"未找到图片 job：{wanted}")
|
||
return {}
|
||
|
||
|
||
def event(state_doc: Dict[str, Any], kind: str, job_id: str, payload: Dict[str, Any]) -> None:
|
||
state_doc.setdefault("events", []).append(
|
||
{
|
||
"time": now_iso(),
|
||
"event": kind,
|
||
"job": job_id,
|
||
**payload,
|
||
}
|
||
)
|
||
|
||
|
||
def infer_origin_out(prompt_path: Path, data: Dict[str, Any]) -> str:
|
||
for key in ("origin_out", "out", "output", "selected_output"):
|
||
value = data.get(key)
|
||
if isinstance(value, str) and value:
|
||
return value
|
||
file_name = data.get("file_name")
|
||
if isinstance(file_name, str) and file_name:
|
||
return str(Path("origin_image") / Path(file_name).name)
|
||
return str(Path("origin_image") / f"{prompt_path.stem}.png")
|
||
|
||
|
||
def infer_asset_out(data: Dict[str, Any], origin_out: str) -> str:
|
||
value = data.get("asset_out")
|
||
if isinstance(value, str) and value:
|
||
return value
|
||
return str(Path("assets") / Path(origin_out).name)
|
||
|
||
|
||
def cmd_init(args: argparse.Namespace) -> int:
|
||
project = project_path(args.project)
|
||
prompts_dir = resolve_in_project(project, args.prompts_dir)
|
||
if not prompts_dir.exists():
|
||
die(f"prompts 目录不存在：{prompts_dir}")
|
||
|
||
jobs: List[Dict[str, Any]] = []
|
||
for prompt_path in sorted(prompts_dir.glob("*.json")):
|
||
data = read_json(prompt_path)
|
||
if not isinstance(data, dict):
|
||
die(f"prompt job 必须是 JSON 对象：{prompt_path}")
|
||
job_id = str(data.get("id") or data.get("job_id") or prompt_path.stem)
|
||
origin_out = infer_origin_out(prompt_path, data)
|
||
jobs.append(
|
||
{
|
||
"id": job_id,
|
||
"prompt_file": str(prompt_path.relative_to(project)),
|
||
"origin_out": origin_out,
|
||
"asset_out": infer_asset_out(data, origin_out),
|
||
"status": "pending",
|
||
"selected_backend": args.selected_backend,
|
||
"created_at": now_iso(),
|
||
}
|
||
)
|
||
|
||
if not jobs:
|
||
die(f"没有找到 prompt JSON：{prompts_dir}")
|
||
|
||
job_doc = {"version": 1, "jobs": jobs}
|
||
state_doc = {"version": 1, "events": [{"time": now_iso(), "event": "init", "job_count": len(jobs)}]}
|
||
save_docs(project, job_doc, state_doc)
|
||
print(json.dumps({"status": "initialized", "jobs": len(jobs)}, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
def cmd_status(args: argparse.Namespace) -> int:
|
||
project = project_path(args.project)
|
||
job_doc = load_job_doc(project)
|
||
counts: Dict[str, int] = {}
|
||
for job in job_doc["jobs"]:
|
||
if not isinstance(job, dict):
|
||
continue
|
||
status = str(job.get("status") or "unknown")
|
||
counts[status] = counts.get(status, 0) + 1
|
||
|
||
print(json.dumps({"project": str(project), "counts": counts}, ensure_ascii=False, indent=2))
|
||
for job in job_doc["jobs"]:
|
||
if isinstance(job, dict):
|
||
print(f"{job_key(job) or '<unknown>'}\t{job.get('status', 'unknown')}\t{job.get('origin_out', '')}")
|
||
return 0
|
||
|
||
|
||
def cmd_dispatch(args: argparse.Namespace) -> int:
|
||
project = project_path(args.project)
|
||
job_doc = load_job_doc(project)
|
||
state_doc = load_state_doc(project)
|
||
job = find_job(job_doc, args.job)
|
||
prompt_file = args.prompt_file or job.get("prompt_file")
|
||
if not isinstance(prompt_file, str) or not prompt_file:
|
||
die("缺少 prompt 文件路径。")
|
||
|
||
job["status"] = "dispatched"
|
||
job["dispatch"] = {
|
||
"agent_id": args.agent_id,
|
||
"agent_nickname": args.agent_nickname,
|
||
"prompt_file": prompt_file,
|
||
"time": now_iso(),
|
||
}
|
||
event(state_doc, "dispatch", args.job, job["dispatch"])
|
||
save_docs(project, job_doc, state_doc)
|
||
print(json.dumps({"status": "dispatched", "job": args.job}, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
def copy_if_needed(source: Path, target: Path) -> None:
|
||
if not source.exists():
|
||
die(f"结果图片不存在：{source}")
|
||
target.parent.mkdir(parents=True, exist_ok=True)
|
||
try:
|
||
if source.resolve() == target.resolve():
|
||
return
|
||
except FileNotFoundError:
|
||
pass
|
||
shutil.copyfile(source, target)
|
||
|
||
|
||
def cmd_result(args: argparse.Namespace) -> int:
|
||
project = project_path(args.project)
|
||
job_doc = load_job_doc(project)
|
||
state_doc = load_state_doc(project)
|
||
job = find_job(job_doc, args.job)
|
||
source = resolve_in_project(project, args.selected_source)
|
||
|
||
origin_value = args.origin_out or job.get("origin_out")
|
||
if not isinstance(origin_value, str) or not origin_value:
|
||
die("缺少 origin_out。")
|
||
origin_out = resolve_in_project(project, origin_value)
|
||
copy_if_needed(source, origin_out)
|
||
|
||
asset_value = args.asset_out or job.get("asset_out")
|
||
if isinstance(asset_value, str) and asset_value:
|
||
copy_if_needed(origin_out, resolve_in_project(project, asset_value))
|
||
|
||
result = {
|
||
"agent_id": args.agent_id,
|
||
"backend_used": args.backend_used,
|
||
"selected_source": str(source),
|
||
"origin_out": str(origin_out.relative_to(project) if origin_out.is_relative_to(project) else origin_out),
|
||
"asset_out": asset_value,
|
||
"qa_note": args.qa_note,
|
||
"time": now_iso(),
|
||
}
|
||
job["status"] = "recorded"
|
||
job["result"] = result
|
||
event(state_doc, "result", args.job, result)
|
||
save_docs(project, job_doc, state_doc)
|
||
print(json.dumps({"status": "recorded", "job": args.job, "origin_out": result["origin_out"]}, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
def cmd_blocker(args: argparse.Namespace) -> int:
|
||
project = project_path(args.project)
|
||
job_doc = load_job_doc(project)
|
||
state_doc = load_state_doc(project)
|
||
job = find_job(job_doc, args.job)
|
||
blocker = {
|
||
"agent_id": args.agent_id,
|
||
"reason": args.reason,
|
||
"evidence": args.evidence,
|
||
"time": now_iso(),
|
||
}
|
||
job["status"] = "blocked"
|
||
job["blocker"] = blocker
|
||
event(state_doc, "blocker", args.job, blocker)
|
||
save_docs(project, job_doc, state_doc)
|
||
print(json.dumps({"status": "blocked", "job": args.job, "reason": args.reason}, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(description="Record MAOMOMO article image job state.")
|
||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||
|
||
init = subparsers.add_parser("init", help="Create slide_jobs.json from prompts/*.json.")
|
||
init.add_argument("project")
|
||
init.add_argument("--prompts-dir", default="prompts")
|
||
init.add_argument("--selected-backend", required=True)
|
||
init.set_defaults(func=cmd_init)
|
||
|
||
status = subparsers.add_parser("status", help="Print job status summary.")
|
||
status.add_argument("project")
|
||
status.set_defaults(func=cmd_status)
|
||
|
||
dispatch = subparsers.add_parser("dispatch", help="Record a dispatched image job.")
|
||
dispatch.add_argument("project")
|
||
dispatch.add_argument("--job", required=True)
|
||
dispatch.add_argument("--agent-id", required=True)
|
||
dispatch.add_argument("--agent-nickname", default="")
|
||
dispatch.add_argument("--prompt-file")
|
||
dispatch.set_defaults(func=cmd_dispatch)
|
||
|
||
result = subparsers.add_parser("result", help="Record a generated image result.")
|
||
result.add_argument("project")
|
||
result.add_argument("--job", required=True)
|
||
result.add_argument("--agent-id", required=True)
|
||
result.add_argument("--backend-used", required=True)
|
||
result.add_argument("--selected-source", required=True)
|
||
result.add_argument("--origin-out")
|
||
result.add_argument("--asset-out")
|
||
result.add_argument("--qa-note", default="")
|
||
result.set_defaults(func=cmd_result)
|
||
|
||
blocker = subparsers.add_parser("blocker", help="Record a blocked image job.")
|
||
blocker.add_argument("project")
|
||
blocker.add_argument("--job", required=True)
|
||
blocker.add_argument("--agent-id", required=True)
|
||
blocker.add_argument("--reason", required=True)
|
||
blocker.add_argument("--evidence", default="")
|
||
blocker.set_defaults(func=cmd_blocker)
|
||
return parser
|
||
|
||
|
||
def main(argv: Optional[Iterable[str]] = None) -> int:
|
||
parser = build_parser()
|
||
args = parser.parse_args(list(argv) if argv is not None else None)
|
||
return args.func(args)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|