555 lines
19 KiB
Python
555 lines
19 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"
|
||
FINAL_MANIFEST_FILE = "final_manifest.json"
|
||
GITHUB_BATCHES_FILE = "github_batches.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 relative_path_string(project: Path, path: Path) -> str:
|
||
try:
|
||
return str(path.resolve().relative_to(project))
|
||
except ValueError:
|
||
return str(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 load_batches_doc(project: Path) -> Dict[str, Any]:
|
||
path = project / GITHUB_BATCHES_FILE
|
||
if not path.exists():
|
||
return {"version": 1, "batches": []}
|
||
raw = read_json(path)
|
||
if not isinstance(raw, dict):
|
||
die(f"{GITHUB_BATCHES_FILE} 必须是对象。")
|
||
if not isinstance(raw.get("batches"), list):
|
||
raw["batches"] = []
|
||
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,
|
||
"aspect_ratio": data.get("aspect_ratio", ""),
|
||
"visual_direction": data.get("visual_direction", ""),
|
||
"delivery_batch": data.get("delivery_batch", ""),
|
||
"qa_status": "pending",
|
||
"entered_assets": False,
|
||
"final": False,
|
||
"github_pushed": False,
|
||
"github_batch": "",
|
||
"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):
|
||
github = "github" if job.get("github_pushed") else "local"
|
||
qa = job.get("qa_status", "")
|
||
print(f"{job_key(job) or '<unknown>'}\t{job.get('status', 'unknown')}\t{qa}\t{github}\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")
|
||
entered_assets = False
|
||
if isinstance(asset_value, str) and asset_value:
|
||
copy_if_needed(origin_out, resolve_in_project(project, asset_value))
|
||
entered_assets = True
|
||
|
||
result = {
|
||
"agent_id": args.agent_id,
|
||
"backend_used": args.backend_used,
|
||
"selected_source": str(source),
|
||
"origin_out": relative_path_string(project, origin_out),
|
||
"asset_out": asset_value,
|
||
"qa_status": "passed",
|
||
"entered_assets": entered_assets,
|
||
"qa_note": args.qa_note,
|
||
"time": now_iso(),
|
||
}
|
||
job["status"] = "recorded"
|
||
job["qa_status"] = "passed"
|
||
job["entered_assets"] = entered_assets
|
||
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 cmd_final_manifest(args: argparse.Namespace) -> int:
|
||
project = project_path(args.project)
|
||
job_doc = load_job_doc(project)
|
||
state_doc = load_state_doc(project)
|
||
assets: List[Dict[str, Any]] = []
|
||
incomplete: List[str] = []
|
||
|
||
for job in job_doc["jobs"]:
|
||
if not isinstance(job, dict):
|
||
continue
|
||
key = job_key(job) or "<unknown>"
|
||
if job.get("status") != "recorded":
|
||
incomplete.append(f"{key}: status={job.get('status', 'unknown')}")
|
||
continue
|
||
if job.get("qa_status") != "passed":
|
||
incomplete.append(f"{key}: qa_status={job.get('qa_status', 'unknown')}")
|
||
continue
|
||
if job.get("entered_assets") is not True:
|
||
incomplete.append(f"{key}: entered_assets=false")
|
||
continue
|
||
|
||
result = job.get("result") if isinstance(job.get("result"), dict) else {}
|
||
asset_value = result.get("asset_out") or job.get("asset_out")
|
||
if not isinstance(asset_value, str) or not asset_value:
|
||
incomplete.append(f"{key}: 缺少 asset_out")
|
||
continue
|
||
|
||
asset_path = resolve_in_project(project, asset_value)
|
||
if not asset_path.exists():
|
||
incomplete.append(f"{key}: assets 文件不存在：{asset_value}")
|
||
continue
|
||
|
||
job["final"] = True
|
||
assets.append(
|
||
{
|
||
"id": key,
|
||
"prompt_file": job.get("prompt_file", ""),
|
||
"origin_out": result.get("origin_out") or job.get("origin_out", ""),
|
||
"asset_out": relative_path_string(project, asset_path),
|
||
"qa_status": job.get("qa_status", ""),
|
||
"qa_note": result.get("qa_note", ""),
|
||
"aspect_ratio": job.get("aspect_ratio", ""),
|
||
"visual_direction": job.get("visual_direction", ""),
|
||
"delivery_batch": job.get("delivery_batch", ""),
|
||
"github_pushed": bool(job.get("github_pushed")),
|
||
"github_batch": job.get("github_batch", ""),
|
||
}
|
||
)
|
||
|
||
if incomplete and not args.allow_incomplete:
|
||
die("仍有未进入最终 manifest 的图片 job：\n" + "\n".join(incomplete))
|
||
if not assets:
|
||
die("没有可写入 final_manifest.json 的最终图片。")
|
||
|
||
manifest_path = resolve_in_project(project, args.out)
|
||
manifest = {
|
||
"version": 1,
|
||
"generated_at": now_iso(),
|
||
"project": project.name,
|
||
"asset_count": len(assets),
|
||
"assets": assets,
|
||
}
|
||
write_json(manifest_path, manifest)
|
||
event(
|
||
state_doc,
|
||
"final_manifest",
|
||
"*",
|
||
{
|
||
"path": relative_path_string(project, manifest_path),
|
||
"asset_count": len(assets),
|
||
"allow_incomplete": args.allow_incomplete,
|
||
},
|
||
)
|
||
save_docs(project, job_doc, state_doc)
|
||
print(
|
||
json.dumps(
|
||
{
|
||
"status": "final_manifest_written",
|
||
"path": relative_path_string(project, manifest_path),
|
||
"assets": len(assets),
|
||
},
|
||
ensure_ascii=False,
|
||
indent=2,
|
||
)
|
||
)
|
||
return 0
|
||
|
||
|
||
def cmd_github_batch(args: argparse.Namespace) -> int:
|
||
project = project_path(args.project)
|
||
job_doc = load_job_doc(project)
|
||
state_doc = load_state_doc(project)
|
||
batches_doc = load_batches_doc(project)
|
||
|
||
rel_paths: List[str] = []
|
||
for value in args.paths:
|
||
path = resolve_in_project(project, value)
|
||
if not path.exists():
|
||
die(f"批次文件不存在：{value}")
|
||
rel_paths.append(relative_path_string(project, path))
|
||
|
||
path_set = set(rel_paths)
|
||
pushed_jobs: List[str] = []
|
||
for job in job_doc["jobs"]:
|
||
if not isinstance(job, dict):
|
||
continue
|
||
result = job.get("result") if isinstance(job.get("result"), dict) else {}
|
||
candidate_values = [job.get("asset_out"), result.get("asset_out")]
|
||
normalized = {
|
||
relative_path_string(project, resolve_in_project(project, value))
|
||
for value in candidate_values
|
||
if isinstance(value, str) and value
|
||
}
|
||
if normalized.intersection(path_set):
|
||
job["github_pushed"] = True
|
||
job["github_batch"] = args.batch
|
||
if isinstance(job.get("result"), dict):
|
||
job["result"]["github_pushed"] = True
|
||
job["result"]["github_batch"] = args.batch
|
||
pushed_jobs.append(job_key(job) or "<unknown>")
|
||
|
||
batch_record = {
|
||
"batch": args.batch,
|
||
"commit": args.commit,
|
||
"url": args.url,
|
||
"release": args.release,
|
||
"paths": rel_paths,
|
||
"pushed_jobs": pushed_jobs,
|
||
"note": args.note,
|
||
"time": now_iso(),
|
||
}
|
||
batches = [item for item in batches_doc["batches"] if not (isinstance(item, dict) and item.get("batch") == args.batch)]
|
||
batches.append(batch_record)
|
||
batches_doc["batches"] = batches
|
||
|
||
manifest_path = project / FINAL_MANIFEST_FILE
|
||
if manifest_path.exists():
|
||
manifest = read_json(manifest_path)
|
||
if isinstance(manifest, dict) and isinstance(manifest.get("assets"), list):
|
||
for item in manifest["assets"]:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
asset_value = item.get("asset_out")
|
||
if not isinstance(asset_value, str) or not asset_value:
|
||
continue
|
||
normalized = relative_path_string(project, resolve_in_project(project, asset_value))
|
||
if normalized in path_set:
|
||
item["github_pushed"] = True
|
||
item["github_batch"] = args.batch
|
||
item["github_commit"] = args.commit
|
||
item["github_url"] = args.url
|
||
write_json(manifest_path, manifest)
|
||
|
||
write_json(project / GITHUB_BATCHES_FILE, batches_doc)
|
||
event(state_doc, "github_batch", "*", batch_record)
|
||
save_docs(project, job_doc, state_doc)
|
||
print(
|
||
json.dumps(
|
||
{
|
||
"status": "github_batch_recorded",
|
||
"batch": args.batch,
|
||
"paths": len(rel_paths),
|
||
"pushed_jobs": pushed_jobs,
|
||
},
|
||
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)
|
||
|
||
final_manifest = subparsers.add_parser("final-manifest", help="Write final_manifest.json from recorded assets.")
|
||
final_manifest.add_argument("project")
|
||
final_manifest.add_argument("--out", default=FINAL_MANIFEST_FILE)
|
||
final_manifest.add_argument("--allow-incomplete", action="store_true")
|
||
final_manifest.set_defaults(func=cmd_final_manifest)
|
||
|
||
github_batch = subparsers.add_parser("github-batch", help="Record a GitHub delivery batch.")
|
||
github_batch.add_argument("project")
|
||
github_batch.add_argument("--batch", required=True)
|
||
github_batch.add_argument("--commit", required=True)
|
||
github_batch.add_argument("--url", required=True)
|
||
github_batch.add_argument("--release", default="")
|
||
github_batch.add_argument("--paths", nargs="+", required=True)
|
||
github_batch.add_argument("--note", default="")
|
||
github_batch.set_defaults(func=cmd_github_batch)
|
||
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())
|