Blame for scripts/local_agent/run_local_agent.py

# /// script
# dependencies = [
#   "torch",
#   "torchvision",
#   "transformers @ git+https://github.com/huggingface/transformers.git",
#   "pillow",
#   "qwen-vl-utils",
#   "bitsandbytes",
#   "accelerate",
# ]
# ///
"""
Single-run, no-conversation grounding agent: loads a capable open-source
vision-language model straight from Hugging Face -- no fine-tuning, no LoRA
adapter, no Anthropic API key anywhere -- and drives it against RULES.md's
protocol for exactly one case (one paper.pdf + one repo/) until it calls
submit_narrative or runs out of turns. Same tool surface, same
schema/narrative.schema.json, same scripts/verify_case.py gate as the
Anthropic-backed harness in scripts/dataset/ and the fine-tuned deployment
script in scripts/finetune/ -- this is the "our app, minus every API call"
version, generalized to accept ANY combination of a local PDF and a
GitHub repo, not just the curated dataset cases already under cases/.

It scaffolds the case directory itself (the same job scripts/new_case.sh
does by hand): drop in a PDF path and a repo (a clone URL, or an
already-cloned/local directory), and it creates cases/<slug>/{paper.pdf,
repo/} before starting the agent loop.

No text-extraction shortcut for the paper (RULES.md is explicit about
this) -- the model has to be able to *read* a rendered page image, so the
model here must be a vision-language ("image-text-to-text") model, not a
plain text LLM. Figure/diagram *judgment calls* still never touch a vision
model -- only the reading of page text does; that discipline lives in the
prompt (RULES.md), not in this driver.

Usage -- local, if you have a CUDA GPU with enough VRAM for --model-id:

    uv run scripts/local_agent/run_local_agent.py \\
        --pdf ~/Downloads/some_paper.pdf \\
        --repo https://github.com/org/repo.git

Usage -- cloud, via HF Jobs (recommended if you don't have a local GPU;
mirrors scripts/finetune/run_finetuned_agent.py's deployment pattern):

    hf jobs uv run --flavor a10g-large --timeout 30m -s HF_TOKEN \\
        -v $(pwd):/workspace:rw \\
        scripts/local_agent/run_local_agent.py -- \\
        --pdf /workspace/inbox/some_paper.pdf \\
        --repo https://github.com/org/repo.git \\
        --case-root /workspace/cases

Pick --model-id to match your hardware. The tool-call parsing here targets
the Qwen chat template's <tool_call>...</tool_call> convention (the same
one scripts/finetune/run_finetuned_agent.py already relies on), so any
Qwen2.5-VL / Qwen3-VL checkpoint is the safe default family; a bigger
sibling (e.g. a 30B+ one) is "more capable" at the cost of more VRAM. A
different model family may use a different tool-call syntax in its chat
template -- check that before pointing --model-id at something else.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import urllib.request
from pathlib import Path

# Overridable via PCD_REPO_ROOT so this exact script can run unmodified
# either from inside a checkout of this repo (the default) or as a
# standalone file uploaded to a remote job (e.g. `hf jobs uv run`) with the
# repo separately mounted read-only at another path (e.g. /workspace) --
# see the HF Jobs usage note below. Doesn't matter for local `uv run`.
_repo_root_override = os.environ.get("PCD_REPO_ROOT")
REPO_ROOT = Path(_repo_root_override) if _repo_root_override else Path(__file__).resolve().parents[2]
RULES_PATH = REPO_ROOT / "RULES.md"
SCHEMA_PATH = REPO_ROOT / "schema" / "narrative.schema.json"
VERIFY_SCRIPT = REPO_ROOT / "scripts" / "verify_case.py"

sys.path.insert(0, str(REPO_ROOT / "scripts" / "dataset"))
from tools import CaseWorkspace, ToolError, TOOL_SCHEMAS, dispatch  # noqa: E402

DEFAULT_MODEL = "Qwen/Qwen3-VL-4B-Instruct"
MAX_TURNS = 40
IMAGE_SIZE = {"shortest_edge": 16384, "longest_edge": 400000}

TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)
URL_RE = re.compile(r"^(https?://|git@|ssh://)")
NARRATIVE_MARKER_START = "===NARRATIVE_JSON_START==="
NARRATIVE_MARKER_END = "===NARRATIVE_JSON_END==="


# ---------------------------------------------------------------- case setup

def slugify(name: str) -> str:
    s = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")
    return re.sub(r"-{2,}", "-", s) or "case"


def setup_case(pdf_source: str, repo_source: str, case_root: Path, slug: str | None) -> tuple[Path, str]:
    """Scaffold cases/<slug>/{paper.pdf, repo/} from an arbitrary PDF (a
    local path or a direct download URL, e.g. an arxiv.org/pdf/... link)
    and an arbitrary repo (clone URL or local directory) -- the headless
    equivalent of scripts/new_case.sh, so this script accepts any paper +
    repo pair, not just an already-prepared case folder. Accepting a PDF
    URL matters for a remote job (e.g. HF Jobs) that has no access to your
    local disk at all -- it can pull both paper and repo over the network
    itself."""
    is_pdf_url = bool(URL_RE.match(pdf_source))
    if not is_pdf_url and not Path(pdf_source).expanduser().is_file():
        raise SystemExit(f"error: --pdf not found: {pdf_source}")

    if is_pdf_url:
        default_stem = pdf_source.rstrip("/").rsplit("/", 1)[-1]
        if default_stem.lower().endswith(".pdf"):
            default_stem = default_stem[:-4]
    else:
        default_stem = Path(pdf_source).stem
    slug = slug or slugify(default_stem)
    case_dir = case_root / slug
    case_dir.mkdir(parents=True, exist_ok=True)

    dest_pdf = case_dir / "paper.pdf"
    if is_pdf_url:
        if not dest_pdf.is_file():
            print(f"downloading {pdf_source} -> {dest_pdf} ...", flush=True)
            urllib.request.urlretrieve(pdf_source, dest_pdf)
    else:
        pdf_path = Path(pdf_source).expanduser()
        if dest_pdf.resolve() != pdf_path.resolve():
            shutil.copy(pdf_path, dest_pdf)

    repo_dir = case_dir / "repo"
    if not repo_dir.exists():
        if URL_RE.match(repo_source) or repo_source.endswith(".git"):
            print(f"cloning {repo_source} -> {repo_dir} ...", flush=True)
            proc = subprocess.run(
                ["git", "clone", "--depth", "1", repo_source, str(repo_dir)],
                capture_output=True, text=True,
            )
            if proc.returncode != 0:
                raise SystemExit(f"error: git clone failed:\n{proc.stderr}")
        else:
            src = Path(repo_source).expanduser()
            if not src.is_dir():
                raise SystemExit(f"error: --repo is neither a git URL nor an existing directory: {repo_source}")
            print(f"copying {src} -> {repo_dir} ...", flush=True)
            shutil.copytree(src, repo_dir)
    else:
        print(f"reusing existing {repo_dir} (already present)", flush=True)

    return case_dir, slug


# ------------------------------------------------------------- system prompt

def build_system_text(paper_title: str | None, venue: str | None) -> str:
    rules = RULES_PATH.read_text()
    schema = SCHEMA_PATH.read_text()

    if paper_title and venue:
        confirmed = (
            f"- Paper title (confirmed, do not re-derive): {paper_title}\n"
            f"- Venue (confirmed, do not re-derive): {venue}\n"
        )
    else:
        confirmed = (
            "- Paper title and venue were NOT provided ahead of time -- read them "
            "off the paper itself (title from the first page; venue from a header/"
            "footer/badge or the arXiv listing, per RULES.md's `case.venue` rule) "
            "and fill them into `case.paper_title` / `case.venue` yourself.\n"
        )

    return f"""{rules}

---

# This run's case

You are grounding exactly one case, headlessly, via tool calls -- no chat,
no user to ask, no second attempt after you call `submit_narrative`
successfully. Work entirely within the tools you've been given: `bash`,
`read_pdf_page`, `read_file`, `crop_figure`, and finally `submit_narrative`
to end the session.

{confirmed}- `paper.pdf` and `repo/` already exist in your working directory -- do not
  try to download or clone anything.
- `read_pdf_page` is the ONLY way to read paper.pdf's content. Do NOT use
  `bash` to run pdftotext, pdftohtml, pdftocairo, grep, strings, or any
  other command against paper.pdf -- none of those give you the actual
  page content, they cannot substitute for `read_pdf_page`, and calling
  them wastes a turn. If you need to see a page, call `read_pdf_page`
  with that page number directly.
- The full JSON Schema narrative.json must satisfy is below. Follow it
  exactly; `submit_narrative` will be schema-checked and then run through
  `scripts/verify_case.py`, the same deterministic gate an interactive
  session uses. If it fails, you'll get the verifier's own output back and
  a chance to fix it -- read the errors literally, they name the exact
  field and problem.
- When genuinely done and confident it would pass, call `submit_narrative`
  once with the complete object. Do not call it speculatively.
- For `code_refs[].snippet`: it is automatically re-read from disk at
  file:start_line-end_line right after you submit and OVERWRITES whatever
  you typed there, so don't spend turns hand-copying it character-for-
  character -- a short placeholder is fine. What has to be exactly right
  is `file`, `start_line`, and `end_line`.

## schema/narrative.schema.json

```json
{schema}
```
"""


def qwen_tools() -> list[dict]:
    return [
        {"type": "function", "function": {
            "name": t["name"], "description": t["description"], "parameters": t["input_schema"],
        }}
        for t in TOOL_SCHEMAS
    ]


def collect_images(messages: list[dict]) -> list:
    imgs = []
    for m in messages:
        content = m.get("content")
        if isinstance(content, list):
            for item in content:
                if isinstance(item, dict) and item.get("type") == "image":
                    imgs.append(item["image"])
    return imgs


def wrap_tool_result(result) -> list[dict]:
    if isinstance(result, list):
        out = []
        for block in result:
            if block.get("type") == "text":
                out.append({"type": "text", "text": block["text"]})
            elif block.get("type") == "image":
                src = block["source"]
                import base64
                import io
                from PIL import Image
                img = Image.open(io.BytesIO(base64.b64decode(src["data"]))).convert("RGB")
                out.append({"type": "image", "image": img})
        return out
    return [{"type": "text", "text": str(result)}]


# ------------------------------------------------------------------- model

def load_model(model_id: str, use_4bit: bool):
    import torch
    from transformers import AutoModelForImageTextToText, AutoProcessor

    cuda_available = torch.cuda.is_available()
    print(f"loading model {model_id} (cuda={cuda_available}, 4bit={use_4bit and cuda_available})...", flush=True)

    processor = AutoProcessor.from_pretrained(model_id, size=IMAGE_SIZE)

    if use_4bit and cuda_available:
        from transformers import BitsAndBytesConfig
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype="bfloat16", bnb_4bit_use_double_quant=True,
        )
        model = AutoModelForImageTextToText.from_pretrained(
            model_id, quantization_config=bnb_config, dtype="bfloat16", device_map="auto",
        )
    else:
        if use_4bit and not cuda_available:
            print("warning: --4bit requires CUDA (bitsandbytes); falling back to full precision", flush=True)
        if cuda_available:
            model = AutoModelForImageTextToText.from_pretrained(model_id, dtype="bfloat16", device_map="auto")
        else:
            mps_available = getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available()
            device = "mps" if mps_available else "cpu"
            # float16 on MPS: roughly half the memory of float32 and well
            # supported by the MPS backend (unlike bfloat16, which still has
            # op-coverage gaps on some torch/macOS versions) -- matters on
            # unified-memory Apple Silicon where headroom is tighter than a
            # dedicated-VRAM GPU. Plain CPU fallback still needs float32.
            dtype = "float16" if device == "mps" else "float32"
            print(f"warning: no CUDA GPU found, running on {device} ({dtype}) -- this will be slow for "
                  f"anything but a small model; consider `hf jobs uv run` for a real GPU instead.", flush=True)
            model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=dtype).to(device)

    model.eval()
    print("model loaded, starting agent loop", flush=True)
    return processor, model


def _dump_narrative(narrative: dict) -> None:
    """Print the full narrative JSON to stdout, wrapped in unique markers.
    This is the retrieval path for a remote job (e.g. HF Jobs) where the
    container's filesystem disappears when the job ends -- `hf jobs logs
    <job_id>` is the only thing guaranteed to survive, so the result has to
    go through it rather than rely solely on the on-disk narrative.json."""
    print(NARRATIVE_MARKER_START, flush=True)
    print(json.dumps(narrative, indent=2, ensure_ascii=False), flush=True)
    print(NARRATIVE_MARKER_END, flush=True)


# ---------------------------------------------------------------------- run

def run_agent(case_dir: Path, processor, model, paper_title: str | None, venue: str | None,
              max_turns: int, max_new_tokens: int) -> bool:
    import torch

    ws = CaseWorkspace(case_dir)
    system_text = build_system_text(paper_title, venue)
    messages: list[dict] = [
        {"role": "system", "content": system_text},
        {"role": "user", "content": [{"type": "text", "text": (
            "Ground the case at this working directory following the protocol and "
            "schema above. Start by reading the PDF's first couple of pages to "
            "find the Method section's boundaries."
        )}]},
    ]
    tools = qwen_tools()
    submitted = None
    last_unverified = None

    for turn in range(1, max_turns + 1):
        text = processor.apply_chat_template(messages, tools=tools, tokenize=False, add_generation_prompt=True)
        images = collect_images(messages)
        inputs = processor(text=[text], images=[images] if images else None, return_tensors="pt").to(model.device)
        with torch.no_grad():
            out_ids = model.generate(
                **inputs, max_new_tokens=max_new_tokens, do_sample=True,
                temperature=0.7, top_p=0.9, repetition_penalty=1.15,
            )
        new_ids = out_ids[0][inputs["input_ids"].shape[1]:]
        generated = processor.tokenizer.decode(new_ids, skip_special_tokens=True)
        print(f"--- turn {turn} ---\n{generated}\n", flush=True)

        messages.append({"role": "assistant", "content": generated})

        calls = TOOL_CALL_RE.findall(generated)
        if not calls:
            messages.append({"role": "user", "content": [{"type": "text", "text": (
                "Continue working the case with tool calls, or call "
                "submit_narrative if you believe it's complete."
            )}]})
            continue

        tool_responses = []
        for raw in calls:
            try:
                call = json.loads(raw)
            except json.JSONDecodeError as e:
                tool_responses.append({"role": "tool", "content": [
                    {"type": "text", "text": f"[error] malformed tool_call JSON: {e}"}
                ]})
                continue
            name, tool_input = call.get("name"), call.get("arguments", {})
            if name == "submit_narrative":
                submitted = tool_input.get("narrative")
                tool_responses.append({"role": "tool", "content": [{"type": "text", "text": "received."}]})
                continue
            try:
                result = dispatch(ws, name, tool_input)
            except ToolError as e:
                result = f"[tool error] {e}"
            tool_responses.append({"role": "tool", "content": wrap_tool_result(result)})
        messages.extend(tool_responses)

        if submitted is None:
            continue

        try:
            resolved = ws.resolve_code_snippets(ws.resolve_crop_placeholders(submitted))
        except ToolError as e:
            messages.append({"role": "user", "content": [{"type": "text", "text": (
                f"submit_narrative referenced a bad crop placeholder: {e}\n\n"
                "Fix the reference (or omit that image) and call submit_narrative again."
            )}]})
            submitted = None
            continue

        narrative_path = case_dir / "narrative.json"
        narrative_path.write_text(json.dumps(resolved, indent=2, ensure_ascii=False) + "\n")
        verify = subprocess.run(
            [sys.executable, str(VERIFY_SCRIPT), str(narrative_path)],
            capture_output=True, text=True,
        )
        print(f"=== verify_case.py (returncode={verify.returncode}) ===\n{verify.stdout}\n{verify.stderr}\n", flush=True)

        if verify.returncode == 0:
            print(f"SUCCESS: {narrative_path} passes verify_case.py", flush=True)
            _dump_narrative(resolved)
            return True

        messages.append({"role": "user", "content": [{"type": "text", "text": (
            f"scripts/verify_case.py rejected that submission:\n\n{verify.stdout}\n{verify.stderr}\n\n"
            "Fix the specific issues named above and call submit_narrative again with the corrected object."
        )}]})
        submitted = None
        last_unverified = resolved

    print(f"FAILED: exceeded max_turns ({max_turns}) without a passing submission", flush=True)
    if last_unverified is not None:
        print("Dumping the last (verifier-rejected) submission below, for partial credit:", flush=True)
        _dump_narrative(last_unverified)
    return False


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--pdf", required=True, help="local path OR direct download URL (e.g. an arxiv.org/pdf/... link) to the paper PDF")
    parser.add_argument("--repo", required=True, help="a git clone URL, or a path to an already-cloned/local repo directory")
    parser.add_argument("--slug", default=None, help="case slug (default: derived from the PDF filename)")
    parser.add_argument("--case-root", default=REPO_ROOT / "cases", type=Path, help="where to scaffold cases/<slug>/ (default: ./cases)")
    parser.add_argument("--paper-title", default=None, help="skip the model having to read this off the PDF itself")
    parser.add_argument("--venue", default=None, help="skip the model having to read this off the PDF itself")
    parser.add_argument("--model-id", default=DEFAULT_MODEL, help=f"any Qwen2.5-VL/Qwen3-VL checkpoint (default: {DEFAULT_MODEL})")
    parser.add_argument("--4bit", dest="use_4bit", action="store_true", help="load in 4-bit via bitsandbytes (CUDA only)")
    parser.add_argument("--max-turns", type=int, default=MAX_TURNS)
    parser.add_argument("--max-new-tokens", type=int, default=4096)
    args = parser.parse_args()

    if not RULES_PATH.is_file():
        raise SystemExit(f"error: {RULES_PATH} not found -- this script must live in the PaperCodeDiagramApp repo")

    # tools.py's read_pdf_page/crop_figure shell out to `pdftoppm` (poppler).
    # Present via Homebrew on a dev Mac already; absent from a bare HF Jobs
    # container image, where apt-get is the only way to get it -- guarded
    # so this is a no-op wherever pdftoppm or apt-get themselves aren't
    # present (e.g. this exact check would otherwise crash on macOS, which
    # has neither reason to run apt-get nor apt-get itself).
    if shutil.which("pdftoppm") is None and shutil.which("apt-get") is not None:
        print("pdftoppm not found -- installing poppler-utils via apt-get...", flush=True)
        subprocess.run(["apt-get", "update"], check=False, capture_output=True)
        subprocess.run(["apt-get", "install", "-y", "poppler-utils"], check=False, capture_output=True)

    case_dir, slug = setup_case(args.pdf, args.repo, args.case_root, args.slug)
    print(f"case ready: {case_dir}", flush=True)

    processor, model = load_model(args.model_id, args.use_4bit)
    ok = run_agent(case_dir, processor, model, args.paper_title, args.venue, args.max_turns, args.max_new_tokens)
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()