seed-data 0.0.1__py3-none-any.whl

doc_gen_agent/__init__.py

@@ -0,0 +1,24 @@
+"""doc-gen-agent — AI-powered document generation pipeline."""
+
+MODELS = {
+    "haiku": {"model_id": "us.anthropic.claude-haiku-4-5-20251001-v1:0", "max_tokens": 63999},
+    "sonnet": {"model_id": "us.anthropic.claude-sonnet-4-6", "max_tokens": 63999},
+    "opus": {"model_id": "us.anthropic.claude-opus-4-6-v1", "max_tokens": 63999},
+    "nova-pro": {"model_id": "us.amazon.nova-pro-v1:0", "max_tokens": 16383},
+    "nova-lite": {"model_id": "us.amazon.nova-lite-v1:0", "max_tokens": 16383},
+    "nova2-lite": {"model_id": "us.amazon.nova-2-lite-v1:0", "max_tokens": 63999},
+    "nova2-pro": {"model_id": "us.amazon.nova-pro-v2:0", "max_tokens": 63999},
+    "nova-premier": {"model_id": "us.amazon.nova-premier-v1:0", "max_tokens": 16383},
+    "deepseek-v3": {"model_id": "deepseek.v3.2", "max_tokens": 16384, "streaming": False},
+    "qwen3-vl": {"model_id": "qwen.qwen3-vl-235b-a22b", "max_tokens": 32767, "streaming": False},
+    "qwen3-coder": {"model_id": "qwen.qwen3-coder-next", "max_tokens": 32767, "streaming": False},
+    "qwen3-coder-480b": {"model_id": "qwen.qwen3-coder-480b-a35b-v1:0", "max_tokens": 32767, "streaming": False},
+    "qwen3-80b": {"model_id": "qwen.qwen3-next-80b-a3b", "max_tokens": 32767, "streaming": False},
+    "gpt-oss": {"model_id": "openai.gpt-oss-120b-1:0", "max_tokens": 16383, "streaming": False},
+    "gpt-oss-20b": {"model_id": "openai.gpt-oss-20b-1:0", "max_tokens": 16383, "streaming": False},
+    "gpt-oss-sg": {"model_id": "openai.gpt-oss-safeguard-120b", "max_tokens": 16383, "streaming": False},
+    "gpt-oss-sg-20b": {"model_id": "openai.gpt-oss-safeguard-20b", "max_tokens": 16383, "streaming": False},
+    "nemotron": {"model_id": "nvidia.nemotron-nano-12b-v2", "max_tokens": 16383, "streaming": False},
+    "nemotron-super": {"model_id": "nvidia.nemotron-super-3-120b", "max_tokens": 16383, "streaming": False},
+    "nemotron-nano-30b": {"model_id": "nvidia.nemotron-nano-3-30b", "max_tokens": 16383, "streaming": False},
+}

doc_gen_agent/__main__.py

@@ -0,0 +1,104 @@
+"""CLI entrypoint: python -m doc_gen_agent"""
+import argparse
+import json
+import sys
+import time
+
+from dotenv import load_dotenv
+
+
+def main():
+    load_dotenv()
+
+    from doc_gen_agent import MODELS
+    from doc_gen_agent.orchestrate import orchestrate
+
+    model_choices = list(MODELS.keys())
+    parser = argparse.ArgumentParser(
+        prog="doc-gen-agent",
+        description="AI-powered document generation pipeline",
+    )
+    parser.add_argument("--schema-dir", required=True, help="Schema directory")
+    parser.add_argument("--output", default="./output", help="Output directory")
+    parser.add_argument("--extra", default=None, help="Extra instructions / brief")
+    parser.add_argument("--count", type=int, default=1, help="Number of docs (>1 uses batch agent)")
+    parser.add_argument("--batch-name", default=None, help="Batch directory name (default: timestamp)")
+    parser.add_argument("--data-model", default="gpt-oss", choices=model_choices)
+    parser.add_argument("--doc-model", default="gpt-oss", choices=model_choices)
+    parser.add_argument("--critic-model", default="sonnet", choices=model_choices)
+    parser.add_argument("--batch-model", default="nova2-lite", choices=model_choices,
+                        help="Model for batch orchestrator")
+    parser.add_argument("--augment", action="store_true")
+    parser.add_argument("--no-critic-samples", action="store_true",
+                        help="Disable reference sample PDFs in doc critic")
+    parser.add_argument("--aug-model", default="deepseek-v3", choices=model_choices)
+    parser.add_argument("--threshold", type=int, default=5)
+    parser.add_argument("--max-attempts", type=int, default=5)
+    parser.add_argument("--timeout", type=int, default=3600)
+    parser.add_argument("--quiet", action="store_true")
+    args = parser.parse_args()
+
+    if args.count > 1:
+        # Batch mode: LLM orchestrator generates diverse documents
+        from doc_gen_agent.batch import run_batch
+        brief = args.extra or "Generate diverse, realistic documents"
+        result = run_batch(
+            schema_dir=args.schema_dir,
+            count=args.count,
+            brief=brief,
+            batch_name=args.batch_name,
+            output_dir=args.output,
+            batch_model=args.batch_model,
+            data_model=args.data_model,
+            doc_model=args.doc_model,
+            critic_model=args.critic_model,
+            aug_model=args.aug_model,
+            threshold=args.threshold,
+            max_attempts=args.max_attempts,
+            timeout=args.timeout,
+            augment=args.augment,
+            critic_samples=not args.no_critic_samples,
+        )
+        print(f"\n{'=' * 60}")
+        print(f"Batch time: {result['elapsed_s']}s")
+    else:
+        # Single doc mode
+        start = time.time()
+        result = orchestrate(
+            schema_dir=args.schema_dir,
+            output_dir=args.output,
+            extra_instructions=args.extra,
+            data_model=args.data_model,
+            doc_model=args.doc_model,
+            critic_model=args.critic_model,
+            aug_model=args.aug_model,
+            threshold=args.threshold,
+            max_attempts=args.max_attempts,
+            timeout=args.timeout,
+            augment=args.augment,
+            verbose=not args.quiet,
+            critic_samples=not args.no_critic_samples,
+        )
+        elapsed = time.time() - start
+
+        print(f"\n{'=' * 60}")
+        print(json.dumps(result, indent=2))
+
+        if result["success"]:
+            print(f"\nPDF:       {result['path']}")
+            if result.get("augmented_path"):
+                print(f"Augmented: {result['augmented_path']}")
+            print(f"Verdict:   {result['verdict']}")
+
+        usage = result.get("token_usage", {})
+        in_tok = usage.get("inputTokens", 0)
+        out_tok = usage.get("outputTokens", 0)
+        cost = (in_tok * 3 + out_tok * 15) / 1_000_000
+        print(f"Tokens:    {in_tok + out_tok:,}  Cost: ${cost:.4f}  Time: {elapsed:.1f}s")
+
+        if not result["success"]:
+            sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

doc_gen_agent/augment.py

@@ -0,0 +1,258 @@
+"""Augraphy augmentation tool wrapper.
+
+Handles PDF → PNG → augraphy pipeline → PNG → PDF conversion.
+The LLM agent decides the augmentation config; this module executes it.
+"""
+import json
+import os
+
+import augraphy
+import numpy as np
+import pypdfium2 as pdfium
+from PIL import Image
+from strands import tool
+
+
+
+# Registry of available augraphy augmentations the agent can choose from
+AUGMENTATION_REGISTRY = {
+    # Ink phase (affects text/ink layer)
+    "InkBleed": augraphy.InkBleed,
+    "InkColorSwap": augraphy.InkColorSwap,
+    "InkMottling": augraphy.InkMottling,
+    "InkShifter": augraphy.InkShifter,
+    "BleedThrough": augraphy.BleedThrough,
+    "Dithering": augraphy.Dithering,
+    "Hollow": augraphy.Hollow,
+    "Letterpress": augraphy.Letterpress,
+    "LinesDegradation": augraphy.LinesDegradation,
+    "LowInkPeriodicLines": augraphy.LowInkPeriodicLines,
+    "LowInkRandomLines": augraphy.LowInkRandomLines,
+    # Paper phase (affects paper/background)
+    "BadPhotoCopy": augraphy.BadPhotoCopy,
+    "ColorPaper": augraphy.ColorPaper,
+    "NoiseTexturize": augraphy.NoiseTexturize,
+    "BrightnessTexturize": augraphy.BrightnessTexturize,
+    "DelaunayTessellation": augraphy.DelaunayTessellation,
+    "VoronoiTessellation": augraphy.VoronoiTessellation,
+    # Post phase (affects final output)
+    "BadPhotoCopy_post": augraphy.BadPhotoCopy,
+    "Brightness": augraphy.Brightness,
+    "ColorShift": augraphy.ColorShift,
+    "DirtyDrum": augraphy.DirtyDrum,
+    "DirtyRollers": augraphy.DirtyRollers,
+    "Faxify": augraphy.Faxify,
+    "Folding": augraphy.Folding,
+    "Gamma": augraphy.Gamma,
+    "Geometric": augraphy.Geometric,
+    "GlitchEffect": augraphy.GlitchEffect,
+    "Jpeg": augraphy.Jpeg,
+    "Markup": augraphy.Markup,
+    "PageBorder": augraphy.PageBorder,
+    "Scribbles": augraphy.Scribbles,
+    "ShadowCast": augraphy.ShadowCast,
+    "SubtleNoise": augraphy.SubtleNoise,
+    "WaterMark": augraphy.WaterMark,
+    "BookBinding": augraphy.BookBinding,
+    "Stains": augraphy.Stains,
+    "DepthSimulatedBlur": augraphy.DepthSimulatedBlur,
+    "LightingGradient": augraphy.LightingGradient,
+    "Moire": augraphy.Moire,
+    "DotMatrix": augraphy.DotMatrix,
+}
+
+
+def _build_augmentation(name: str, params: dict):
+    """Instantiate an augraphy augmentation by name with given params."""
+    cls = AUGMENTATION_REGISTRY.get(name)
+    if cls is None:
+        raise ValueError(f"Unknown augmentation: {name}. Available: {list(AUGMENTATION_REGISTRY.keys())}")
+    # Convert list params to tuples, and ensure integer params stay as ints
+    # (JSON parse turns everything to float)
+    converted = {}
+    for k, v in params.items():
+        if isinstance(v, list) and len(v) == 2 and all(isinstance(x, (int, float)) for x in v):
+            # If both values are whole numbers, convert to int tuple
+            if all(x == int(x) for x in v):
+                converted[k] = (int(v[0]), int(v[1]))
+            else:
+                converted[k] = (v[0], v[1])
+        elif isinstance(v, float) and v == int(v):
+            converted[k] = int(v)
+        else:
+            converted[k] = v
+    return cls(**converted)
+
+
+def _pdf_to_images(pdf_path: str, dpi: int = 200) -> list[np.ndarray]:
+    """Convert PDF pages to numpy arrays."""
+    doc = pdfium.PdfDocument(pdf_path)
+    images = []
+    for page in doc:
+        bitmap = page.render(scale=dpi / 72)
+        img = np.asarray(bitmap.to_pil().convert("RGB"))
+        images.append(img)
+        page.close()
+    doc.close()
+    return images
+
+
+def _images_to_pdf(images: list[np.ndarray], output_path: str, dpi: int = 200):
+    """Convert numpy arrays back to a PDF."""
+    pages = [Image.fromarray(img) for img in images]
+    # resolution embeds DPI so the physical page size matches the source render
+    pages[0].save(
+        output_path,
+        format="PDF",
+        save_all=True,
+        append_images=pages[1:],
+        resolution=float(dpi),
+    )
+
+
+def run_augmentation(pdf_path: str, config: dict, dpi: int = 150, timeout: int = 90) -> str:
+    """Apply augraphy augmentations to a PDF.
+
+    Args:
+        pdf_path: Path to the input PDF.
+        config: Dict with keys "ink_phase", "paper_phase", "post_phase".
+                Each is a list of {"name": str, "params": dict}.
+        dpi: Resolution for PDF rendering (lower = faster).
+        timeout: Max seconds per page augmentation.
+
+    Returns:
+        Path to the augmented PDF.
+    """
+
+    ink_phase = [_build_augmentation(a["name"], a.get("params", {}))
+                 for a in config.get("ink_phase", [])]
+    paper_phase = [_build_augmentation(a["name"], a.get("params", {}))
+                   for a in config.get("paper_phase", [])]
+    post_phase = [_build_augmentation(a["name"], a.get("params", {}))
+                  for a in config.get("post_phase", [])]
+
+    pipeline = augraphy.AugraphyPipeline(
+        ink_phase=ink_phase,
+        paper_phase=paper_phase,
+        post_phase=post_phase,
+    )
+
+    images = _pdf_to_images(pdf_path, dpi=dpi)
+
+    augmented_images = []
+    for i, img in enumerate(images):
+        augmented = pipeline(img)
+        augmented_images.append(augmented)
+        print(f"  Page {i+1}: augmented successfully")
+        # Explicitly free original image to reduce peak memory
+        del img
+
+    images.clear()
+
+    base, ext = os.path.splitext(pdf_path)
+    output_path = f"{base}_augmented.pdf"
+    _images_to_pdf(augmented_images, output_path, dpi=dpi)
+
+    # Explicitly free augmented images after writing
+    augmented_images.clear()
+    import gc
+    gc.collect()
+
+    # Clean up augraphy cache
+    import shutil as _shutil
+    cache_dir = os.path.join(os.getcwd(), "augraphy_cache")
+    if os.path.isdir(cache_dir):
+        _shutil.rmtree(cache_dir, ignore_errors=True)
+
+    return output_path
+
+
+@tool
+def apply_augmentation(pdf_path: str, augmentation_config: str) -> dict:
+    """Apply augraphy augmentations to a PDF document.
+
+    Takes a clean PDF and applies document aging/degradation effects.
+
+    Args:
+        pdf_path: Path to the clean PDF to augment.
+        augmentation_config: JSON string with augmentation configuration.
+            Must have keys "ink_phase", "paper_phase", "post_phase".
+            Each is a list of objects with "name" and "params".
+            ONLY use these safe augmentations:
+            - ink_phase: InkBleed, LowInkRandomLines, LowInkPeriodicLines, LinesDegradation, Dithering
+            - paper_phase: NoiseTexturize, BrightnessTexturize, ColorPaper
+            - post_phase: SubtleNoise, Jpeg, Brightness, Gamma, DirtyDrum, DirtyRollers, Faxify, BadPhotoCopy
+
+    Returns:
+        Path to the augmented PDF, or error message.
+    """
+    try:
+        config = json.loads(augmentation_config)
+
+        # Log the augmentation chain
+        all_augs = []
+        for phase in ("ink_phase", "paper_phase", "post_phase"):
+            for aug in config.get(phase, []):
+                all_augs.append(f"  [{phase}] {aug['name']}: {aug.get('params', {})}")
+        chain_log = "\n".join(all_augs) if all_augs else "  (none)"
+        print(f"\nAugmentation chain:\n{chain_log}")
+
+        output_path = run_augmentation(pdf_path, config)
+        if os.path.exists(output_path):
+            size = os.path.getsize(output_path)
+            # Save the aug config for reproducibility
+            config_path = output_path.replace("_augmented.pdf", "_aug_config.json")
+            with open(config_path, "w") as f:
+                json.dump(config, f, indent=2)
+            return {
+                "status": "success",
+                "content": [{"text": f"Augmented PDF saved to: {output_path} ({size:,} bytes)\nConfig saved to: {config_path}"}],
+            }
+        else:
+            return {
+                "status": "error",
+                "content": [{"text": "Augmentation completed but output file was not created."}],
+            }
+    except Exception as e:
+        return {
+            "status": "error",
+            "content": [{"text": f"Augmentation failed: {e}. Try simpler augmentations (SubtleNoise, Jpeg, Brightness)."}],
+        }
+
+
+def critique_augmented_document(pdf_path: str, model: str = "haiku", threshold: int = 7) -> dict:
+    """Critique an augmented PDF using structured_output."""
+    from doc_gen_agent import prompts as p
+    from doc_gen_agent.critique import CritiqueResult
+    from doc_gen_agent.utils import make_model as _make_model
+    from strands import Agent
+
+    with open(pdf_path, "rb") as f:
+        pdf_bytes = f.read()
+
+    system_prompt = p.render("aug_critic", threshold=threshold)
+    agent = Agent(model=_make_model(model), system_prompt=system_prompt)
+
+    user_message = [
+        {"document": {"format": "pdf", "name": "augmented_doc",
+                       "source": {"bytes": pdf_bytes}}},
+        {"text": f"Evaluate this augmented document. Threshold: {threshold}/10."},
+    ]
+
+    result = agent(user_message, structured_output_model=CritiqueResult)
+    critique = result.structured_output
+    verdict = "accepted" if critique.score >= threshold else "rejected"
+
+    print(f"\n--- Aug Critique ({model}) ---")
+    print(f"  Score:   {critique.score}/10 → {verdict}")
+    print(f"  Summary: {critique.summary}")
+    for issue in critique.issues:
+        print(f"  [{issue.severity}] ({issue.category}) {issue.description}")
+    print("--- End Aug Critique ---\n")
+
+    return {
+        "score": critique.score,
+        "verdict": verdict,
+        "issues": [i.model_dump() for i in critique.issues],
+        "summary": critique.summary,
+    }

doc_gen_agent/batch.py

@@ -0,0 +1,213 @@
+"""Batch orchestrator — LLM agent that generates diverse documents using orchestrate as a tool."""
+import glob
+import json
+import os
+import shutil
+import sys
+import time
+from datetime import datetime
+
+from strands import Agent, tool
+
+from doc_gen_agent.orchestrate import orchestrate
+from doc_gen_agent.utils import make_model
+
+_pipeline_config = {}
+_batch_results = []
+
+
+@tool
+def generate_document(extra_instructions: str) -> dict:
+    """Generate a single document using the full pipeline.
+
+    Args:
+        extra_instructions: Specific scenario for this document. Be detailed:
+            include advertiser name, industry, station/location, time period,
+            and any other specifics that make this document unique.
+
+    Returns:
+        Summary of the generation result.
+    """
+    cfg = _pipeline_config
+    start = time.time()
+
+    result = orchestrate(
+        schema_dir=cfg["schema_dir"],
+        output_dir=cfg["batch_dir"],
+        extra_instructions=extra_instructions,
+        data_model=cfg["data_model"],
+        doc_model=cfg["doc_model"],
+        critic_model=cfg["critic_model"],
+        aug_model=cfg["aug_model"],
+        threshold=cfg["threshold"],
+        max_attempts=cfg["max_attempts"],
+        timeout=cfg["timeout"],
+        augment=cfg["augment"],
+        critic_samples=cfg.get("critic_samples", True),
+        verbose=False,
+    )
+    elapsed = time.time() - start
+
+    doc_id = result.get("doc_id", "unknown")
+
+    # Track result
+    usage = result.get("token_usage", {})
+    entry = {
+        "doc_id": doc_id,
+        "success": result["success"],
+        "verdict": result.get("verdict", "error"),
+        "extra": extra_instructions,
+        "pdf": result.get("path"),
+        "data_json": result.get("data_json"),
+        "augmented": result.get("augmented_path"),
+        "time_s": round(elapsed, 1),
+        "tokens": usage.get("inputTokens", 0) + usage.get("outputTokens", 0),
+    }
+    if not result["success"]:
+        entry["error"] = result.get("error", "unknown")
+    _batch_results.append(entry)
+
+    status = "✓" if result["success"] else "✗"
+    print(f"  {status} [{len(_batch_results)}] {extra_instructions[:60]}... ({elapsed:.0f}s)")
+    if not result["success"]:
+        print(f"    ERROR: {result.get('error', 'unknown')[:120]}")
+
+    # Bail early if first 2 docs both fail (likely a config/auth issue)
+    if len(_batch_results) >= 2 and all(not r["success"] for r in _batch_results):
+        msg = "First 2 documents failed — likely a configuration issue. Aborting batch."
+        print(f"\n  ⚠ {msg}")
+        return {"status": "error", "content": [{"text": msg}]}
+
+    # Return concise summary to keep batch agent context small
+    summary = {k: entry[k] for k in ("success", "verdict", "doc_id", "time_s")}
+    return {"status": "success", "content": [{"text": json.dumps(summary)}]}
+
+
+def run_batch(
+    schema_dir: str,
+    count: int,
+    brief: str,
+    batch_name: str = None,
+    output_dir: str = "./output",
+    batch_model: str = "sonnet",
+    data_model: str = "sonnet",
+    doc_model: str = "sonnet",
+    critic_model: str = "haiku",
+    aug_model: str = "sonnet",
+    threshold: int = 7,
+    max_attempts: int = 5,
+    timeout: int = 3600,
+    augment: bool = False,
+    critic_samples: bool = True,
+) -> dict:
+    """Run the batch orchestrator agent.
+
+    Creates a batch directory under output_dir with subdirectories for
+    pdfs/, data/, and augmented/. Writes a batch_manifest.json at the end.
+    """
+    global _pipeline_config, _batch_results
+    _batch_results = []
+
+    # Create batch directory
+    if not batch_name:
+        batch_name = datetime.now().strftime("batch_%Y%m%d_%H%M%S")
+    batch_dir = os.path.join(output_dir, batch_name)
+    os.makedirs(batch_dir, exist_ok=True)
+
+    # Copy schema and steering docs into config/ folder for reproducibility
+    config_dir = os.path.join(batch_dir, "config")
+    os.makedirs(config_dir, exist_ok=True)
+    for f in glob.glob(os.path.join(schema_dir, "*.json")) + glob.glob(os.path.join(schema_dir, "*.md")):
+        shutil.copy2(f, config_dir)
+
+    _pipeline_config = {
+        "schema_dir": schema_dir,
+        "batch_dir": batch_dir,
+        "data_model": data_model,
+        "doc_model": doc_model,
+        "critic_model": critic_model,
+        "aug_model": aug_model,
+        "threshold": threshold,
+        "max_attempts": max_attempts,
+        "timeout": timeout,
+        "augment": augment,
+        "critic_samples": critic_samples,
+    }
+
+    system_prompt = f"""You are a batch document generation orchestrator. Your job is to
+generate {count} diverse, unique documents by calling generate_document repeatedly.
+
+RULES:
+- Call generate_document exactly {count} times
+- Each call MUST have completely different extra_instructions
+- Vary: advertiser names, industries, station locations, time periods, spot types
+- Track what you've generated to avoid repetition
+- If a generation fails, note it and continue with the next one
+- After all {count} calls, output a brief summary line: "BATCH COMPLETE: X/Y succeeded"
+- Do NOT explain your plan — just start generating immediately
+"""
+
+    agent = Agent(
+        model=make_model(batch_model),
+        system_prompt=system_prompt,
+        tools=[generate_document],
+    )
+
+    print(f"Batch:  {batch_name}")
+    print(f"Count:  {count}")
+    print(f"Brief:  {brief}")
+    print(f"Output: {batch_dir}")
+    print("=" * 60)
+
+    batch_start = time.time()
+    agent(f"Generate {count} diverse documents. Brief: {brief}")
+    batch_elapsed = time.time() - batch_start
+
+    # Write manifest
+    succeeded = sum(1 for r in _batch_results if r["success"])
+    total_tokens = sum(r.get("tokens", 0) for r in _batch_results)
+    total_cost = total_tokens * 3 / 1_000_000  # rough estimate
+
+    # Get git hash for reproducibility
+    import subprocess as _sp
+    try:
+        git_hash = _sp.check_output(
+            ["git", "rev-parse", "HEAD"], stderr=_sp.DEVNULL, text=True
+        ).strip()
+    except Exception:
+        git_hash = "unknown"
+
+    manifest = {
+        "batch_name": batch_name,
+        "brief": brief,
+        "command": " ".join(sys.argv),
+        "git_hash": git_hash,
+        "count_requested": count,
+        "count_succeeded": succeeded,
+        "count_failed": len(_batch_results) - succeeded,
+        "total_tokens": total_tokens,
+        "est_cost_usd": round(total_cost, 4),
+        "elapsed_s": round(batch_elapsed, 1),
+        "config": {
+            "schema_dir": schema_dir,
+            "data_model": data_model,
+            "doc_model": doc_model,
+            "critic_model": critic_model,
+            "aug_model": aug_model,
+            "threshold": threshold,
+            "augment": augment,
+        },
+        "documents": _batch_results,
+    }
+    manifest_path = os.path.join(batch_dir, "config", "batch_manifest.json")
+    os.makedirs(os.path.dirname(manifest_path), exist_ok=True)
+    with open(manifest_path, "w") as f:
+        json.dump(manifest, f, indent=2)
+
+    print(f"\n{'=' * 60}")
+    print(f"Batch complete: {succeeded}/{count} succeeded")
+    print(f"Manifest: {manifest_path}")
+    print(f"Tokens:   {total_tokens:,}  Est cost: ${total_cost:.4f}")
+    print(f"Time:     {batch_elapsed:.1f}s ({batch_elapsed/max(count,1):.1f}s avg)")
+
+    return manifest

doc_gen_agent/cli.py

@@ -0,0 +1,27 @@
+"""Shared CLI argument parser for doc-gen-agent scripts."""
+import argparse
+from doc_gen_agent import MODELS
+
+
+def base_parser(description="doc-gen-agent"):
+    """Create the base argument parser with common options."""
+    model_choices = list(MODELS.keys())
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument("--schema-dir", required=True, help="Schema directory")
+    parser.add_argument("--output", default="./output", help="Output directory")
+    parser.add_argument("--extra", default=None, help="Extra instructions / brief")
+    parser.add_argument("--data-model", default="gpt-oss", choices=model_choices)
+    parser.add_argument("--doc-model", default="nova2-lite", choices=model_choices)
+    parser.add_argument("--critic-model", default="sonnet", choices=model_choices)
+    parser.add_argument("--augment", action="store_true")
+    parser.add_argument("--aug-model", default="gpt-oss", choices=model_choices)
+    parser.add_argument("--no-critic-samples", action="store_true",
+                        help="Disable reference sample PDFs in doc critic")
+    parser.add_argument("--renderer", default="weasyprint", choices=["reportlab", "weasyprint"],
+                        help="PDF rendering engine (default: weasyprint)")
+    parser.add_argument("--preview", action="store_true",
+                        help="Enable preview_pdf tool for visual self-inspection (requires VL model)")
+    parser.add_argument("--threshold", type=int, default=5)
+    parser.add_argument("--max-attempts", type=int, default=5)
+    parser.add_argument("--timeout", type=int, default=3600)
+    return parser