310 lines
11 KiB
Python
Executable File
310 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""auto-draft — schedule ai-pipeline work based on spare subscription capacity.
|
|
|
|
OUTLINE REVISION. This revision is dryrun-only: it probes the subscription, computes
|
|
the spare capacity, and prints the three priority queues. Job execution is not yet
|
|
implemented. See specs/auto-draft.spec.md for the full design.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
|
|
# --- Constants --------------------------------------------------------------
|
|
|
|
VALID_PROVIDERS = ("anthropic", "minimax")
|
|
WEEKLY_PERIOD_DAYS = 7
|
|
WEEKLY_PERIOD_SECONDS = WEEKLY_PERIOD_DAYS * 86400
|
|
|
|
# Inline HTML-comment scoring convention (see spec).
|
|
SCORE_RE = re.compile(r"<!--\s*score:\s*(\d+)\s*-->")
|
|
EXEMPT_RE = re.compile(r"<!--\s*exempt:\s*true\s*-->")
|
|
TITLE_RE = re.compile(r"^(\d+)-(.+)$")
|
|
|
|
# --- ANSI colour ------------------------------------------------------------
|
|
|
|
_use_colour = sys.stdout.isatty()
|
|
_BOLD = "\033[1m" if _use_colour else ""
|
|
_DIM = "\033[2m" if _use_colour else ""
|
|
_GREEN = "\033[32m" if _use_colour else ""
|
|
_YELLOW = "\033[33m" if _use_colour else ""
|
|
_RED = "\033[31m" if _use_colour else ""
|
|
_CYAN = "\033[36m" if _use_colour else ""
|
|
_RESET = "\033[0m" if _use_colour else ""
|
|
|
|
|
|
# --- Helpers ----------------------------------------------------------------
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def _fmt_duration(seconds: float) -> str:
|
|
seconds = int(seconds)
|
|
if seconds < 0:
|
|
seconds = 0
|
|
days, rem = divmod(seconds, 86400)
|
|
hours, rem = divmod(rem, 3600)
|
|
minutes, _ = divmod(rem, 60)
|
|
if days > 0:
|
|
return f"{days}d {hours}h"
|
|
if hours > 0:
|
|
return f"{hours}h {minutes}m"
|
|
return f"{minutes}m"
|
|
|
|
|
|
def _fmt_pct(pct: float) -> str:
|
|
return f"{pct:.1f}%"
|
|
|
|
|
|
def _fmt_pct_signed(pct: float) -> str:
|
|
sign = "+" if pct >= 0 else ""
|
|
return f"{sign}{pct:.1f}%"
|
|
|
|
|
|
# --- Provider probe ---------------------------------------------------------
|
|
|
|
def probe_provider(provider: str) -> dict | None:
|
|
"""Run agent-subscriptions --output json and return the matching provider entry."""
|
|
try:
|
|
result = subprocess.run(
|
|
["agent-subscriptions", "--output", "json"],
|
|
capture_output=True, text=True, check=True, timeout=30,
|
|
)
|
|
except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError) as e:
|
|
print(f"{_YELLOW}Warning: agent-subscriptions probe failed: {e}{_RESET}", file=sys.stderr)
|
|
return None
|
|
|
|
try:
|
|
data = json.loads(result.stdout)
|
|
except json.JSONDecodeError as e:
|
|
print(f"{_YELLOW}Warning: agent-subscriptions returned invalid JSON: {e}{_RESET}", file=sys.stderr)
|
|
return None
|
|
|
|
for p in data.get("providers", []):
|
|
if p.get("provider", "").lower() == provider.lower():
|
|
return p
|
|
return None
|
|
|
|
|
|
# --- Spare capacity ---------------------------------------------------------
|
|
|
|
def compute_spare(provider_record: dict | None, threshold_pct: float) -> dict | None:
|
|
"""Compute the linear limit and spare capacity for the seven-day window."""
|
|
if not provider_record or not provider_record.get("available"):
|
|
return None
|
|
seven_day = provider_record.get("windows", {}).get("seven_day", {})
|
|
reset_at_str = seven_day.get("reset_at")
|
|
usage_pct = seven_day.get("utilization_pct")
|
|
if not reset_at_str or usage_pct is None:
|
|
return None
|
|
|
|
reset_at = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00"))
|
|
renewed_at = reset_at - timedelta(seconds=WEEKLY_PERIOD_SECONDS)
|
|
now = _now()
|
|
elapsed_seconds = max(0.0, (now - renewed_at).total_seconds())
|
|
elapsed_pct = (elapsed_seconds / WEEKLY_PERIOD_SECONDS) * 100.0
|
|
limit_pct = (elapsed_seconds / WEEKLY_PERIOD_SECONDS) * threshold_pct
|
|
spare_pct = limit_pct - usage_pct
|
|
|
|
return {
|
|
"reset_at": reset_at,
|
|
"renewed_at": renewed_at,
|
|
"elapsed_seconds": elapsed_seconds,
|
|
"elapsed_pct": elapsed_pct,
|
|
"threshold_pct": threshold_pct,
|
|
"limit_pct": limit_pct,
|
|
"usage_pct": usage_pct,
|
|
"spare_pct": spare_pct,
|
|
}
|
|
|
|
|
|
def print_spare_capacity(provider: str, spare: dict | None) -> None:
|
|
if spare is None:
|
|
print(f"{_YELLOW}[auto-draft] No spare capacity data for {provider} (provider unavailable or reset_at missing){_RESET}")
|
|
return
|
|
|
|
print(f"{_BOLD}[auto-draft] Spare capacity for {provider}{_RESET}")
|
|
print(f" Reset at: {spare['reset_at'].isoformat()}")
|
|
print(f" Renewed at: {spare['renewed_at'].isoformat()}")
|
|
print(f" Elapsed: {_fmt_duration(spare['elapsed_seconds'])} of {WEEKLY_PERIOD_DAYS}d ({spare['elapsed_pct']:.1f}%)")
|
|
print(f" Threshold: {_fmt_pct(spare['threshold_pct'])}")
|
|
print(f" Linear limit: {_fmt_pct(spare['limit_pct'])}")
|
|
print(f" Current usage: {_fmt_pct(spare['usage_pct'])}")
|
|
|
|
spare_pct = spare["spare_pct"]
|
|
if spare_pct >= 0:
|
|
colour = _GREEN
|
|
summary = "under budget"
|
|
else:
|
|
colour = _RED
|
|
summary = "over budget; do not queue any work"
|
|
print(f" Spare capacity: {colour}{_fmt_pct_signed(spare_pct)}{_RESET} ({summary})")
|
|
print()
|
|
|
|
|
|
# --- Title discovery --------------------------------------------------------
|
|
|
|
def discover_titles(base_dir: Path) -> list[tuple[int, str, str, Path]]:
|
|
"""Return list of (title_id, title_slug, original_stem, parent_dir) sorted by title_id."""
|
|
titles: list[tuple[int, str, str, Path]] = []
|
|
seen: set[tuple[str, Path]] = set()
|
|
for f in sorted(base_dir.rglob("*.overview.md")):
|
|
stem = f.name.removesuffix(".overview.md")
|
|
key = (stem, f.parent)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
m = TITLE_RE.match(stem)
|
|
if m:
|
|
title_id = int(m.group(1))
|
|
title_slug = m.group(2)
|
|
else:
|
|
title_id = float("inf")
|
|
title_slug = stem
|
|
titles.append((title_id, title_slug, stem, f.parent))
|
|
titles.sort(key=lambda x: (x[0], x[1]))
|
|
return titles
|
|
|
|
|
|
def _files_for(parent: Path, stem: str, provider: str) -> tuple[Path, Path]:
|
|
score = parent / f"{stem}.ai_score.{provider}.md"
|
|
generated = parent / f"{stem}.ai_generated.{provider}.md"
|
|
return score, generated
|
|
|
|
|
|
# --- Score parsing ----------------------------------------------------------
|
|
|
|
def parse_score_file(score_path: Path, polish_score: int) -> tuple[int, int]:
|
|
"""Return (paragraphs_below, total_scored)."""
|
|
try:
|
|
content = score_path.read_text()
|
|
except OSError:
|
|
return 0, 0
|
|
paragraphs = re.split(r"\n\s*\n", content)
|
|
below = 0
|
|
scored = 0
|
|
for para in paragraphs:
|
|
if not para.strip():
|
|
continue
|
|
score_m = SCORE_RE.search(para)
|
|
if not score_m:
|
|
continue
|
|
scored += 1
|
|
score = int(score_m.group(1))
|
|
is_exempt = bool(EXEMPT_RE.search(para))
|
|
if score < polish_score and not is_exempt:
|
|
below += 1
|
|
return below, scored
|
|
|
|
|
|
# --- Queue building ---------------------------------------------------------
|
|
|
|
def build_queues(base_dir: Path, provider: str, polish_score: int):
|
|
"""Return (polish_queue, scoring_queue, generation_queue). Each queue is a list of dicts."""
|
|
polish_q: list[dict] = []
|
|
scoring_q: list[dict] = []
|
|
generation_q: list[dict] = []
|
|
|
|
for title_id, title_slug, stem, parent in discover_titles(base_dir):
|
|
score_file, generated_file = _files_for(parent, stem, provider)
|
|
if score_file.exists():
|
|
below, total = parse_score_file(score_file, polish_score)
|
|
if below > 0:
|
|
polish_q.append({
|
|
"title_id": title_id,
|
|
"title_slug": title_slug,
|
|
"stem": stem,
|
|
"parent": parent,
|
|
"below": below,
|
|
"total_scored": total,
|
|
})
|
|
elif generated_file.exists():
|
|
scoring_q.append({
|
|
"title_id": title_id,
|
|
"title_slug": title_slug,
|
|
"stem": stem,
|
|
"parent": parent,
|
|
})
|
|
else:
|
|
generation_q.append({
|
|
"title_id": title_id,
|
|
"title_slug": title_slug,
|
|
"stem": stem,
|
|
"parent": parent,
|
|
})
|
|
|
|
return polish_q, scoring_q, generation_q
|
|
|
|
|
|
def _sort_queue(q: list[dict]) -> list[dict]:
|
|
return sorted(q, key=lambda x: (x["title_id"], x["title_slug"]))
|
|
|
|
|
|
def print_queue(label: str, kind: str, queue: list[dict]) -> None:
|
|
queue = _sort_queue(queue)
|
|
plural = "s" if len(queue) != 1 else ""
|
|
print(f"{_BOLD}[auto-draft] {label} ({kind}) — {len(queue)} item{plural}{_RESET}")
|
|
if not queue:
|
|
print(f" {_DIM}(empty){_RESET}")
|
|
else:
|
|
for item in queue:
|
|
stem = item["stem"]
|
|
if kind == "polish":
|
|
print(f" {stem} {item['below']} paragraph(s) below score (of {item['total_scored']} scored)")
|
|
else:
|
|
print(f" {stem}")
|
|
print()
|
|
|
|
|
|
# --- Main -------------------------------------------------------------------
|
|
|
|
def parse_args(argv: list[str]) -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(
|
|
prog="auto-draft",
|
|
description="Schedule ai-pipeline work based on spare subscription capacity (outline revision).",
|
|
)
|
|
p.add_argument("--provider", required=True, choices=VALID_PROVIDERS,
|
|
help="Subscription provider to draw capacity from")
|
|
p.add_argument("--threshold", required=True, type=float,
|
|
help="Maximum subscription utilization over the period, as a percentage")
|
|
p.add_argument("--polish-score", required=True, type=int,
|
|
help="Paragraph score below which a paragraph is queued for polishing")
|
|
p.add_argument("--base-dir", default=".", type=Path,
|
|
help="Root directory to scan for dossier files (default: current directory)")
|
|
p.add_argument("-n", "--dryrun", action="store_true", default=True,
|
|
help="Show capacity and queues without executing (default: dryrun)")
|
|
args = p.parse_args(argv)
|
|
if not (0 < args.threshold <= 100):
|
|
p.error("--threshold must be between 0 and 100 (exclusive of 0)")
|
|
if not (0 <= args.polish_score <= 10):
|
|
p.error("--polish-score must be between 0 and 10")
|
|
return args
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
args = parse_args(argv if argv is not None else sys.argv[1:])
|
|
base_dir = args.base_dir.expanduser().resolve()
|
|
if not base_dir.is_dir():
|
|
print(f"Error: --base-dir {base_dir} is not a directory", file=sys.stderr)
|
|
return 1
|
|
|
|
provider_record = probe_provider(args.provider)
|
|
spare = compute_spare(provider_record, args.threshold)
|
|
print_spare_capacity(args.provider, spare)
|
|
|
|
polish_q, scoring_q, generation_q = build_queues(base_dir, args.provider, args.polish_score)
|
|
print_queue("Polish queue", "polish", polish_q)
|
|
print_queue("Scoring queue", "scoring", scoring_q)
|
|
print_queue("Generation queue", "generation", generation_q)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|