coorl-lost-cities/scripts/librarian_survey.py

"""Stage 2 survey mode: promote-scan every archive without a research counterpart.

Walks `docs/archive/*.md`, filters out entries that already have a
`docs/research/<stem>.md` (after stripping any -YYYY-MM-DD suffix),
and dispatches each remaining archive to the configured LLM via the
same prompt assembly as `scripts/librarian_promote.py`. Outputs
land under `runs/tmp/librarian-survey-<timestamp>/`, one file per
archive.

Per archive output:
- `<stem>.md` — markdown draft, when the LLM judged the entry
  promotable (ready to copy into `docs/research/`).
- `<stem>.SKIP.txt` — the LLM's one-line reason, when it judged
  the entry non-promotable.
- `<stem>.ERROR.txt` — captured stderr, when the CLI itself failed.

Sequential by default; pass `--parallel N` to dispatch up to N concurrent
LLM calls (subprocess.run is mostly waiting on the network, so threads
are enough — no GIL fight).

Usage:
    uv run python scripts/librarian_survey.py
    LIBRARIAN_LLM=gemini uv run python scripts/librarian_survey.py
    uv run python scripts/librarian_survey.py --dry-run    # list candidates only
    uv run python scripts/librarian_survey.py --max 3      # cap candidates
    uv run python scripts/librarian_survey.py --parallel 4 # 4 concurrent calls
"""

from __future__ import annotations

import argparse
import os
import re
import subprocess
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path

LLM_COMMANDS = {
    "claude": ["claude", "-p"],
    "codex": ["codex", "exec"],
    "gemini": ["gemini", "-p", ""],
}

DATE_SUFFIX = re.compile(r"-\d{4}-\d{2}-\d{2}$")
LAST_VERIFIED = re.compile(r"^\*\*Last verified:\*\*[^\n]*$", re.MULTILINE)


def _repo_root() -> Path:
    current = Path(__file__).resolve()
    for parent in current.parents:
        if (parent / "pyproject.toml").is_file():
            return parent
    raise RuntimeError("pyproject.toml을 찾을 수 없어 repository root를 판정할 수 없습니다.")


def _assemble_prompt(
    system_prompt: str,
    rel_archive: Path,
    archive_body: str,
    rel_target: Path,
) -> str:
    return (
        f"{system_prompt}\n\n"
        "---\n\n"
        "Task: Draft a `docs/research/` note from the archive entry below.\n"
        "Follow the rules in your system prompt above (style template, "
        "`Last verified:` and `Source:` headers, `file:line` citations "
        "verified against the current tree, ~1 page, prose over bullet "
        "soup).\n\n"
        f"**Source archive:** `{rel_archive}`\n"
        f"**Suggested target filename:** `{rel_target}`\n\n"
        "If the archive does not contain a durable conclusion (e.g. it "
        "is a one-off bench result with no general lesson), respond "
        "with a single line `SKIP: <reason>` instead of a draft.\n\n"
        "Output: the markdown content of the new file only. No "
        "preamble, no code fences around the whole thing, no "
        "explanation after. Begin with the H1 header line.\n\n"
        "---\n\n"
        "Archive body:\n\n"
        f"{archive_body}\n"
    )


def _current_commit_sha(root: Path) -> str:
    """Resolve HEAD to a short SHA. Returns 'unknown' if git is unavailable."""
    result = subprocess.run(
        ["git", "rev-parse", "--short", "HEAD"],
        capture_output=True,
        text=True,
        check=False,
        cwd=root,
    )
    if result.returncode != 0:
        return "unknown"
    return result.stdout.strip() or "unknown"


def _post_process_draft(content: str, commit_sha: str) -> str:
    """Normalize a draft before writing to disk.

    The LLM is unreliable about running `git` itself in headless mode; it
    tends to leave `<short-hash>` placeholders or literal `HEAD` strings.
    Rewrite the `**Last verified:**` line in place with today's date and
    the real short SHA. If the line is missing entirely, leave the draft
    untouched so the omission stays visible during review.
    """
    today = datetime.now().strftime("%Y-%m-%d")
    new_line = f"**Last verified:** {today}, commit `{commit_sha}`"
    return LAST_VERIFIED.sub(new_line, content, count=1)


def _has_counterpart(archive_stem_no_date: str, research_stems: set[str]) -> bool:
    """True if a research note already covers this archive entry.

    Checks exact match plus tail match: archive `deep-cfr-foo-bar`
    is considered covered if research `foo-bar` exists, since some
    research notes intentionally drop a domain prefix.
    """
    if archive_stem_no_date in research_stems:
        return True
    parts = archive_stem_no_date.split("-")
    for i in range(1, len(parts)):
        tail = "-".join(parts[i:])
        if tail in research_stems:
            return True
    return False


def _dispatch_one(
    archive: Path,
    target: Path,
    *,
    root: Path,
    cmd: list[str],
    system_prompt: str,
    commit_sha: str,
) -> tuple[Path, str, str, str]:
    """Process a single archive end-to-end.

    Returns ``(rel_archive, status, stem, payload)`` where ``status`` is
    one of ``"draft"``, ``"skip"``, or ``"error"``. ``payload`` is the
    text to write under ``out_dir / f"{stem}.<ext>"``. Pure function
    (apart from the subprocess call) so it is safe to call from a thread.
    """
    rel_archive = archive.relative_to(root)
    rel_target = target.relative_to(root)
    archive_body = archive.read_text(encoding="utf-8")
    prompt = _assemble_prompt(system_prompt, rel_archive, archive_body, rel_target)
    result = subprocess.run(
        cmd,
        input=prompt,
        capture_output=True,
        text=True,
        check=False,
    )
    stem_out = DATE_SUFFIX.sub("", archive.stem)
    if result.returncode != 0:
        return (rel_archive, "error", stem_out, result.stderr)
    output = result.stdout.strip()
    if output.lstrip().upper().startswith("SKIP"):
        return (rel_archive, "skip", stem_out, output)
    return (rel_archive, "draft", stem_out, _post_process_draft(output, commit_sha))


def _candidates(root: Path) -> list[tuple[Path, Path]]:
    """Return (archive_path, suggested_research_target) for archives lacking a counterpart."""
    archive_dir = root / "docs" / "archive"
    research_dir = root / "docs" / "research"
    research_stems = {p.stem for p in research_dir.glob("*.md")}
    out: list[tuple[Path, Path]] = []
    for path in sorted(archive_dir.glob("*.md")):
        stem = DATE_SUFFIX.sub("", path.stem)
        if _has_counterpart(stem, research_stems):
            continue
        target = research_dir / f"{stem}.md"
        out.append((path, target))
    return out


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="List candidate archives and their suggested research targets; do not call the LLM.",
    )
    parser.add_argument(
        "--max",
        type=int,
        default=None,
        help="Process at most N archives (testing/cost guard).",
    )
    parser.add_argument(
        "--parallel",
        type=int,
        default=1,
        help=(
            "Number of concurrent LLM calls. Default 1 (sequential). "
            "Try 4 for a meaningful speedup on large surveys."
        ),
    )
    args = parser.parse_args()
    if args.parallel < 1:
        parser.error("--parallel must be >= 1")

    root = _repo_root()
    candidates = _candidates(root)
    if args.max is not None:
        candidates = candidates[: args.max]

    if not candidates:
        print("No archive entries without a research counterpart. Nothing to do.")
        return 0

    if args.dry_run:
        print(f"{len(candidates)} archive entries lack a research counterpart:")
        for archive, target in candidates:
            print(f"  {archive.relative_to(root)}  →  {target.relative_to(root)}")
        return 0

    backend = os.environ.get("LIBRARIAN_LLM", "claude").lower()
    cmd = LLM_COMMANDS.get(backend)
    if cmd is None:
        print(
            f"Unknown LIBRARIAN_LLM={backend}; supported: {sorted(LLM_COMMANDS)}",
            file=sys.stderr,
        )
        return 1

    system_prompt = (root / "scripts" / "librarian-prompt.md").read_text(encoding="utf-8")
    commit_sha = _current_commit_sha(root)

    timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
    out_dir = root / "runs" / "tmp" / f"librarian-survey-{timestamp}"
    out_dir.mkdir(parents=True, exist_ok=True)

    drafts = 0
    skips = 0
    errors = 0
    total = len(candidates)
    print_lock = threading.Lock()

    def _record(idx: int, rel_archive: Path, status: str, stem: str, payload: str) -> str:
        if status == "error":
            (out_dir / f"{stem}.ERROR.txt").write_text(payload, encoding="utf-8")
            note = "ERROR"
        elif status == "skip":
            (out_dir / f"{stem}.SKIP.txt").write_text(payload, encoding="utf-8")
            note = f"SKIP ({payload[:80]})"
        else:
            (out_dir / f"{stem}.md").write_text(payload, encoding="utf-8")
            note = f"draft ({len(payload)} chars)"
        with print_lock:
            print(f"[{idx}/{total}] {rel_archive} → {backend}", file=sys.stderr)
            print(f"    {note}", file=sys.stderr)
        return status

    if args.parallel == 1:
        for idx, (archive, target) in enumerate(candidates, 1):
            rel_archive, status, stem_out, payload = _dispatch_one(
                archive,
                target,
                root=root,
                cmd=cmd,
                system_prompt=system_prompt,
                commit_sha=commit_sha,
            )
            kind = _record(idx, rel_archive, status, stem_out, payload)
            drafts += kind == "draft"
            skips += kind == "skip"
            errors += kind == "error"
    else:
        with ThreadPoolExecutor(max_workers=args.parallel) as pool:
            futures = {
                pool.submit(
                    _dispatch_one,
                    archive,
                    target,
                    root=root,
                    cmd=cmd,
                    system_prompt=system_prompt,
                    commit_sha=commit_sha,
                ): None
                for archive, target in candidates
            }
            for idx, fut in enumerate(as_completed(futures), 1):
                rel_archive, status, stem_out, payload = fut.result()
                kind = _record(idx, rel_archive, status, stem_out, payload)
                drafts += kind == "draft"
                skips += kind == "skip"
                errors += kind == "error"

    print()
    print(f"Survey complete: {drafts} drafts, {skips} skips, {errors} errors.")
    print(f"Output: {out_dir.relative_to(root)}")
    if drafts:
        print()
        print("Next: review drafts, then for each acceptable one:")
        print(f"  cp {out_dir.relative_to(root)}/<name>.md docs/research/<name>.md")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())