scripts/ts03-material-index

#!/usr/bin/env python3
"""ts03-material-index — build per-subject study-material snapshots from Google Drive.

Run this on a schedule (cron/launchd), separate from the bot process. It uses the
OWNER's full-tool Claude Code session (with the Google Drive MCP connector added) to
list files per subject folder, then extracts each file ONE AT A TIME — large textbooks
and short slide decks in the same folder are handled independently, so one huge PDF
timing out doesn't lose the other files' data. Writes one JSON file per subject under
~/.cache/ts03-material/<Folder>.json.

The bot's member sandbox (zero tools, see ts03-bot) only ever reads these local JSON
files — it never touches Google Drive directly. That keeps the existing isolation
guarantee (member tier = zero tools) intact even though members can now get answers
grounded in real course material.

Usage:
  ./scripts/ts03-material-index            # index every folder in the routes file
  ./scripts/ts03-material-index GPU        # index only the named folder(s)
  ./scripts/ts03-material-index GPU NLP

Setup:
  1. Add the Google Drive MCP server to Claude Code (same way you added Sheets MCP).
  2. Subject folders live at the top level of My Drive: GPU/, NLP/, AI/, DL/, ...
  3. Copy templates/subject-routes.json.example to ~/.config/ts03-subject-routes.json
     and edit the alias -> folder-name map to match your actual folder names.
  4. Run this script once by hand to confirm it works, then schedule it.
"""
import json
import os
import re
import shutil
import subprocess
import sys
import time

ROOT_FOLDER = os.environ.get("TS03_DRIVE_ROOT", "")   # "" = subject folders at top level of My Drive
MATERIAL_DIR = os.path.expanduser("~/.cache/ts03-material")
# claude runs HERE, not in the repo: it downloads book PDFs and writes text copies into its working
# directory, and those must never land in a git checkout. Wiped at the end of every run.
WORK_DIR = os.path.join(MATERIAL_DIR, "_work")
ROUTES_FILE = os.path.expanduser(os.environ.get("SUBJECT_ROUTES_FILE", "~/.config/ts03-subject-routes.json"))
MAX_CHARS_PER_FILE = int(os.environ.get("MATERIAL_MAX_CHARS_PER_FILE", "12000"))
LIST_TIMEOUT = int(os.environ.get("MATERIAL_LIST_TIMEOUT", "120"))
FILE_TIMEOUT = int(os.environ.get("MATERIAL_FILE_TIMEOUT", "600"))   # per file — big PDFs need room
CLAUDE_BIN = os.environ.get("CLAUDE_BIN", "claude")
# Earlier this forced "sonnet" explicitly, to avoid a since-fixed bug where a retired model
# snapshot ID 404'd in freshly-spawned shells. That's now fixed at the root (stale PATH
# ordering), and direct testing showed sonnet is measurably LESS reliable than this user's
# own configured default at the multi-step "list folders, match, use that ID" Drive task --
# every forced-sonnet attempt failed where the default model succeeded. Don't override unless
# explicitly asked to: inherit whatever's already confirmed working for this install.
INDEX_MODEL = os.environ.get("INDEX_MODEL", "")


def run(prompt, timeout):
    args = [CLAUDE_BIN] + (["--model", INDEX_MODEL] if INDEX_MODEL else []) + ["-p", prompt]
    os.makedirs(WORK_DIR, exist_ok=True)
    r = subprocess.run(args, capture_output=True, text=True, timeout=timeout, cwd=WORK_DIR)
    if r.returncode != 0:
        print(f"claude exited {r.returncode}: {(r.stderr or '')[:300]}", file=sys.stderr)
    return (r.stdout or "").strip()


def load_routes():
    try:
        with open(ROUTES_FILE) as f:
            return json.load(f)
    except FileNotFoundError:
        sys.exit(f"Routes file not found: {ROUTES_FILE}\n"
                 f"Copy templates/subject-routes.json.example there and edit it first.")
    except json.JSONDecodeError as e:
        sys.exit(f"Routes file is not valid JSON: {e}")


def all_folders():
    return sorted(set(load_routes().values()))


def where_folder(folder):
    # Drive may store the folder with a trailing slash (e.g. "GPU/" not "GPU"). A name-based
    # exact-match search query legitimately returns zero against that. Don't hint and hope --
    # spell out the exact procedure that's confirmed to work: list top-level folders, match
    # ignoring a trailing slash, then use THAT folder's ID for everything after.
    steps = (f'To find it: FIRST list the top-level folders of My Drive directly (not a name '
             f'search query). Find the one whose name, with any trailing slash stripped, equals '
             f'"{folder}". Use THAT folder\'s ID for every call below -- do not search by name.')
    return (f'the folder at path "{ROOT_FOLDER}/{folder}". {steps}' if ROOT_FOLDER
            else f'the top-level Drive folder named "{folder}". {steps}')


LIST_RETRIES = int(os.environ.get("MATERIAL_LIST_RETRIES", "2"))   # extra attempts if it comes back empty


def list_files(folder):
    """Lightweight call: just the file names, no content yet.

    An empty result here is genuinely ambiguous: it's confirmed, by direct observation, that
    an IDENTICAL prompt against a non-empty Drive folder can either correctly walk the Drive
    exploration steps or shortcut straight to an empty answer, depending on the run -- this
    isn't a bug to hunt further, it's agentic tool-calling non-determinism. Retry before
    trusting a zero result, rather than accepting the first answer as ground truth."""
    prompt = f"""Using the Google Drive tools, list the files directly inside {where_folder(folder)}
in My Drive (not Shared Drives, unless told otherwise). Only include pdf, ppt, pptx, doc, docx
files. Reply with ONLY one JSON object, no prose, no code fences:
{{"files": ["exact file name 1", "exact file name 2"]}}
If the folder doesn't exist or has no matching files, reply {{"files": []}}."""

    for attempt in range(1 + LIST_RETRIES):
        files = _list_files_once(folder, prompt)
        if files:
            return files
        if attempt < LIST_RETRIES:
            print(f"[{folder}] empty result on attempt {attempt + 1}, retrying...", file=sys.stderr)
    return []


def _list_files_once(folder, prompt):
    try:
        out = run(prompt, LIST_TIMEOUT)
    except subprocess.TimeoutExpired:
        print(f"[{folder}] listing timed out after {LIST_TIMEOUT}s", file=sys.stderr)
        return []
    m = re.search(r"\{.*\}", out, re.S)
    if not m:
        print(f"[{folder}] could not list files: {out[:200]!r}", file=sys.stderr)
        return []
    try:
        return json.loads(m.group(0)).get("files", [])
    except json.JSONDecodeError as e:
        print(f"[{folder}] bad list JSON: {e}", file=sys.stderr)
        return []


def extract_file(folder, name):
    """One file per call — a 40MB textbook and a 500KB slide deck are independent calls,
    so a slow/failing book never costs you the rest of the folder."""
    prompt = f"""Using the Google Drive tools, open the file "{name}" inside {where_folder(folder)}
in My Drive, and read its content (slide text for ppt/pptx, body text for pdf/doc).

PDFs and slide decks are NOT text/* MIME type, so a direct text-read tool will refuse them —
that is expected, not a failure. If that happens, convert the file to a Google Doc first
(e.g. a Drive "convert to Google Docs" / export tool), then read THAT. For .ppt/.pptx slide
decks, if a Google Doc conversion is not available, use whatever conversion or export the Drive
tools offer for slide decks (for example Google Slides) and read the slide text instead.
Do this automatically, do not stop and ask for confirmation.

IMPORTANT: if you create a temporary converted copy (a Google Doc or Google Slides file) to read
the content, DELETE that copy immediately after you've read it, before replying. This runs on a daily schedule --
leaving the copy behind would create a new duplicate file in the user's Drive every single day.
Only delete the copy YOU just created for this conversion, never the original PDF/PPT.

It may be a large textbook — do not try to reproduce it in full. Pick out only the most
exam-relevant parts: definitions, key formulas, algorithms, worked examples, chapter or
section summaries. Skip covers, indexes, references, boilerplate and repeated headers/footers.

Reply with up to {MAX_CHARS_PER_FILE} characters of PLAIN TEXT only — no JSON, no code
fences, no preamble, no "here is the extracted text".

If, after converting, the document is genuinely empty or near-empty (this usually means the
PDF is scanned page images with no OCR text layer), reply with exactly: EMPTY (scanned, no text layer)
For any other failure, reply with exactly: EMPTY (<short reason>)"""
    try:
        out = run(prompt, FILE_TIMEOUT).strip()
    except subprocess.TimeoutExpired:
        print(f"[{folder}] {name}: timed out after {FILE_TIMEOUT}s", file=sys.stderr)
        return ""
    if not out or out.startswith("EMPTY"):
        reason = out[5:].strip(" ()") if out.startswith("EMPTY") and len(out) > 5 else "no reason given"
        print(f"[{folder}] {name}: nothing extracted ({reason})", file=sys.stderr)
        return ""
    return out[:MAX_CHARS_PER_FILE]


def load_cache(folder):
    try:
        with open(os.path.join(MATERIAL_DIR, folder + ".json")) as f:
            return json.load(f)
    except (OSError, ValueError):
        return {}


def index_subject(folder):
    # Never let a bad run wipe text that is already indexed: the model has returned an EMPTY list
    # for a non-empty folder before, and a single file can fail to extract on a given run.
    prev = {f["name"]: f["text"] for f in load_cache(folder).get("files", []) if f.get("name") and f.get("text")}
    names = list_files(folder)
    if not names:
        print(f"[{folder}] no files found, keeping the previous cache ({len(prev)} file(s))")
        return True
    files = []
    for name in names:
        text = extract_file(folder, name)
        if text:
            status = f"ok, {len(text)} chars"
        elif name in prev:
            text, status = prev[name], f"extraction failed, kept previous ({len(prev[name])} chars)"
        else:
            status = "skipped"
        print(f"[{folder}] {name}: {status}")
        if text:
            files.append({"name": name, "text": text})

    data = {"files": files, "indexed_at": time.strftime("%Y-%m-%d %H:%M")}
    os.makedirs(MATERIAL_DIR, exist_ok=True)
    tmp = os.path.join(MATERIAL_DIR, folder + ".json.tmp")
    with open(tmp, "w") as f:
        json.dump(data, f)
    os.replace(tmp, os.path.join(MATERIAL_DIR, folder + ".json"))
    print(f"[{folder}] indexed {len(files)}/{len(names)} file(s)")
    return True   # a folder with 0 files isn't a failure — just nothing to cache yet


def main():
    requested = sys.argv[1:]
    folders = requested if requested else all_folders()
    unknown = set(requested) - set(all_folders())
    if unknown:
        print(f"Not in {ROUTES_FILE}: {', '.join(sorted(unknown))}", file=sys.stderr)
    ok = True
    try:
        for folder in folders:
            ok = index_subject(folder) and ok
    finally:
        shutil.rmtree(WORK_DIR, ignore_errors=True)
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()