Protect plan.md from bad Gemma replies

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

1 files changed, 62 insertions, 16 deletions

diff --git a/fanfictioner b/fanfictioner
index 34ae021..c993761 100755
--- a/fanfictioner
+++ b/fanfictioner
@@ -16,6 +16,7 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
 import argparse
+import http.client
 import json
 import os
 import random
@@ -282,6 +283,8 @@ def check_json(pj, plan):
     if sorted(p["n"] for p in pj["pages"]) != sorted(want):
         errs.append("pages differ from plan.md")
     for p in pj["pages"]:
+        if p["orientation"] not in ("portrait", "landscape"):
+            errs.append(f"page {p['n']}: bad orientation {p['orientation']}")
         got, exp = [e["characters"] for e in p["edits"]], edit_groups(want.get(p["n"], []))
         if got != exp:
             errs.append(f"page {p['n']}: edit groups {got}, expected {exp}")
@@ -397,16 +400,16 @@ def review(d, profile, prompt, size, refs=(), skip=False):
                 sys.exit("quit, run again to resume")
 
 
-def llm(path, body=None):
+def llm(path, body=None, timeout=900):
     data = json.dumps(body).encode() if body is not None else None
     req = urllib.request.Request(LLM_URL + path, data=data, headers={"Content-Type": "application/json"})
-    with urllib.request.urlopen(req, timeout=900) as r:
+    with urllib.request.urlopen(req, timeout=timeout) as r:
         return json.load(r)
 
 
 def check_llm():
     try:
-        ids = [m["id"] for m in llm("/v1/models")["data"]]
+        ids = [m["id"] for m in llm("/v1/models", timeout=10)["data"]]
     except OSError as e:
         sys.exit(f"llama-server unreachable at {LLM_URL}: {e}")
     if LLM_MODEL not in ids:
@@ -416,7 +419,7 @@ def check_llm():
 def unload_llm():
     """Free VRAM for sd-cli. Idempotent: an unloaded or unreachable server is fine."""
     try:
-        llm("/models/unload", {"model": LLM_MODEL})
+        llm("/models/unload", {"model": LLM_MODEL}, timeout=10)
     except OSError:
         pass
 
@@ -437,8 +440,11 @@ def strip_fences(s):
 def compile_plan(md, plan):
     """plan.md -> plan.json dict, or None when Gemma twice returns groups that do not match."""
     for _ in range(2):
-        pj = json.loads(chat(COMPILE_SYS, f"{md}\n\nEdit groups, in order:\n{edit_hint(plan)}", SCHEMA))
-        errs = check_json(pj, plan)
+        try:
+            pj = json.loads(chat(COMPILE_SYS, f"{md}\n\nEdit groups, in order:\n{edit_hint(plan)}", SCHEMA))
+            errs = check_json(pj, plan)
+        except (ValueError, KeyError, TypeError) as e:
+            errs = [f"bad JSON from Gemma: {e}"]
         if not errs:
             break
         print("compile mismatch:", *errs, sep="\n  ")
@@ -451,6 +457,9 @@ def compile_plan(md, plan):
     return pj
 
 
+GEMMA_ERRS = (OSError, ValueError, KeyError, http.client.HTTPException)
+
+
 def plan_stage(story, wd):
     md_path, js_path = wd / "plan.md", wd / "plan.json"
     if js_path.exists() and md_path.exists() and js_path.stat().st_mtime >= md_path.stat().st_mtime:
@@ -458,7 +467,11 @@ def plan_stage(story, wd):
     check_llm()
     wd.mkdir(exist_ok=True)
     if not md_path.exists():
-        md_path.write_text(strip_fences(chat(DRAFT_SYS, story.read_text())) + "\n")
+        try:
+            draft = strip_fences(chat(DRAFT_SYS, story.read_text())) + "\n"
+        except GEMMA_ERRS as e:
+            sys.exit(f"Gemma failed: {e}")
+        md_path.write_text(draft)
     while True:
         md = md_path.read_text()
         plan = parse_plan(md)
@@ -466,21 +479,51 @@ def plan_stage(story, wd):
         print(f"\n{md}\n--- {md_path}: {len(plan['characters'])} characters, {len(plan['pages'])} pages")
         for e in errs:
             print("  !", e)
-        k = input("[c]onfirm [e]dit [g]emma revise [q]uit > ").strip().lower()
-        if k == "c" and not errs:
-            pj = compile_plan(md, plan)
-            if pj:
-                js_path.write_text(json.dumps(pj, indent=2) + "\n")
-                return pj
+        try:
+            k = input("[c]onfirm [e]dit [g]emma revise [q]uit > ").strip().lower()
+        except EOFError:
+            k = "q"
+        if k == "c" and errs:
+            print("fix the ! lines first")
+        elif k == "c":
+            try:
+                pj = compile_plan(md, plan)
+            except GEMMA_ERRS as e:
+                print(f"Gemma failed: {e}")
+                continue
+            if not pj:
+                continue
+            if md_path.read_text() != md:
+                print("plan.md changed during compile, review again")
+                continue
+            tmp = js_path.with_name("plan.json.tmp")
+            tmp.write_text(json.dumps(pj, indent=2) + "\n")
+            os.replace(tmp, js_path)
+            return pj
         elif k == "e":
             editor(md_path)
         elif k == "g":
-            ins = input("instruction for Gemma: ")
-            md_path.write_text(strip_fences(chat(
-                DRAFT_SYS, f"Current plan:\n\n{md}\n\nRevise it: {ins}\nReturn the complete new plan.")) + "\n")
+            try:
+                ins = input("instruction for Gemma: ")
+            except EOFError:
+                continue
+            try:
+                new = strip_fences(chat(
+                    DRAFT_SYS, f"Current plan:\n\n{md}\n\nRevise it: {ins}\nReturn the complete new plan.")) + "\n"
+            except GEMMA_ERRS as e:
+                print(f"Gemma failed: {e}")
+                continue
+            nerrs = validate_plan(parse_plan(new))
+            if nerrs:
+                (wd / "plan.rejected.md").write_text(new)
+                print("revision invalid, plan.md kept, see plan.rejected.md", *nerrs, sep="\n  ")
+            else:
+                (wd / "plan.md.bak").write_text(md)
+                md_path.write_text(new)
         elif k == "q":
             sys.exit(0)
 
+
 def selftest():
     plan = parse_plan(SAMPLE_PLAN)
     assert (plan["title"], plan["series"], plan["book"]) == ("Rooftop", "Test Series", "Book One"), plan
@@ -544,6 +587,9 @@ def selftest():
                                       {"characters": ["Jon", "Kit"], "prompt": "x"}]
     assert check_json(regrouped, plan) == [
         "page 1: edit groups [['Mara'], ['Jon', 'Kit']], expected [['Mara', 'Jon'], ['Kit']]"]
+    badori = json.loads(json.dumps(good))
+    badori["pages"][1]["orientation"] = "square"
+    assert check_json(badori, plan) == ["page 2: bad orientation square"]
     short = json.loads(json.dumps(good))
     del short["characters"][2], short["pages"][1]
     assert check_json(short, plan) == ["characters differ from plan.md", "pages differ from plan.md"]