user-docs

/**
 * RAG Python venv setup utility.
 *
 * Pattern: ensure the RAG venv lazily and repair it when validation fails.
 *
 * - Venv lives at <pluginRoot>/.rag-venv/
 * - Ready marker: .rag-venv/rag_ready.marker — written once after a successful
 *   install so subsequent starts skip pip entirely (fast path).
 * - If the marker exists but the venv python is gone, the marker is removed and
 *   setup runs again.
 * - onStatus callback receives human-readable progress lines → forward to
 *   ctx.status() in tool implementations.
 */

import fs from "fs";
import os from "os";
import path from "path";
import { spawn, execFileSync } from "child_process";

const IS_WIN = process.platform === "win32";

// ── Paths ─────────────────────────────────────────────────────────────────────

export function pluginRoot(): string {
  return process.cwd();
}

export function ragVenvDir(): string {
  return path.join(pluginRoot(), ".rag-venv");
}

export function ragVenvPython(): string {
  return IS_WIN
    ? path.join(ragVenvDir(), "Scripts", "python.exe")
    : path.join(ragVenvDir(), "bin", "python3");
}

function ragVenvPip(): string {
  return IS_WIN
    ? path.join(ragVenvDir(), "Scripts", "pip.exe")
    : path.join(ragVenvDir(), "bin", "pip");
}

function ragReadyMarker(): string {
  return path.join(ragVenvDir(), "rag_ready.marker");
}

function requirementsFile(): string {
  return path.join(pluginRoot(), "python", "requirements.txt");
}

function logsDir(): string {
  return path.join(pluginRoot(), "logs");
}

// ── Helpers ───────────────────────────────────────────────────────────────────

function ragVenvExists(): boolean {
  return fs.existsSync(ragVenvPython());
}

function pythonImportsWork(pythonPath: string, imports: string[]): boolean {
  if (!fs.existsSync(pythonPath)) return false;
  try {
    execFileSync(
      pythonPath,
      ["-c", imports.map((name) => `import ${name}`).join("; ")],
      { stdio: "ignore", timeout: 10_000 }
    );
    return true;
  } catch {
    return false;
  }
}

function pythonImportsWorkLogged(
  pythonPath: string,
  imports: string[],
  logFile: string
): boolean {
  if (!fs.existsSync(pythonPath)) {
    try { fs.appendFileSync(logFile, `[import-check] python not found: ${pythonPath}\n`); } catch {}
    return false;
  }
  try {
    execFileSync(
      pythonPath,
      ["-c", imports.map((name) => `import ${name}`).join("; ")],
      { stdio: "ignore", timeout: 10_000 }
    );
    return true;
  } catch (err: any) {
    // Re-run with stderr captured so we can log the actual import error.
    try {
      execFileSync(
        pythonPath,
        ["-c", imports.map((name) => `import ${name}`).join("; ")],
        { encoding: "utf-8", timeout: 10_000 }
      );
    } catch (inner: any) {
      const detail = (inner?.stderr || inner?.stdout || String(inner)).slice(0, 2000);
      try { fs.appendFileSync(logFile, `[import-check] FAILED for ${imports.join(", ")}:\n${detail}\n`); } catch {}
    }
    return false;
  }
}

function unlinkIfExists(filePath: string): void {
  try { if (fs.existsSync(filePath)) fs.unlinkSync(filePath); } catch {}
}

function ragVenvPythonVersion(): string | null {
  try {
    const out = execFileSync(ragVenvPython(), ["--version"], {
      encoding: "utf-8",
      timeout: 5_000,
    });
    const m = (out || "").trim().match(/Python\s+(\d+\.\d+)/);
    return m ? m[1] : null;
  } catch {
    return null;
  }
}

// Stream a subprocess line-by-line; resolve on exit code 0, reject otherwise.
function runAndStream(
  cmd: string,
  args: string[],
  cwd: string,
  onLine: (line: string) => void,
  logFile?: string
): Promise<void> {
  return new Promise((resolve, reject) => {
    const child = spawn(cmd, args, { cwd, stdio: ["ignore", "pipe", "pipe"] });
    const logStream = logFile
      ? fs.createWriteStream(logFile, { flags: "a" })
      : null;

    function onData(data: Buffer) {
      for (const line of data.toString().split("\n")) {
        const t = line.trim();
        if (t) {
          onLine(t);
          logStream?.write(t + "\n");
        }
      }
    }

    child.stdout.on("data", onData);
    child.stderr.on("data", onData);
    child.on("error", (err) => {
      logStream?.end();
      reject(err);
    });
    child.on("close", (code) => {
      logStream?.end();
      if (code === 0) resolve();
      else reject(new Error(`${path.basename(cmd)} exited with code ${code}`));
    });
  });
}

// ── Python discovery ──────────────────────────────────────────────────────────

/**
 * Returns absolute paths to Python executables bundled inside LM Studio itself
 * (the _amphibian vendor directory). Sorted newest-first by directory name.
 */
function lmsAmphibianPythonCandidates(): string[] {
  const amphibianDir = path.join(
    os.homedir(), ".lmstudio", "extensions", "backends", "vendor", "_amphibian"
  );
  if (!fs.existsSync(amphibianDir)) return [];
  try {
    const entries = fs.readdirSync(amphibianDir)
      .filter(e => /^cpython3\.\d+-(win|mac|linux)/i.test(e))
      .sort()
      .reverse(); // highest version / build number first
    return entries.flatMap(entry => {
      const base = path.join(amphibianDir, entry);
      return IS_WIN
        ? [
            path.join(base, "bin", "python3.exe"),
            path.join(base, "bin", "python.exe"),
            path.join(base, "python.exe"),
          ]
        : [
            path.join(base, "bin", "python3"),
            path.join(base, "bin", "python"),
          ];
    });
  } catch {
    return [];
  }
}

function systemPythonCandidates(minMajor: number, minMinor: number): string[] {
  const candidates: string[] = [];

  // NOTE: LM Studio's bundled Python (_amphibian) is intentionally NOT listed here.
  // It is compiled with macOS hardened runtime (codesign flags=runtime), which prevents
  // it from loading third-party native extensions (e.g. pymupdf/_extra.so) that have a
  // different Team ID. A venv created from that Python inherits the restriction and the
  // import check silently fails. On Windows there is no Team ID enforcement, so it would
  // work there — but we still want consistent behaviour across platforms.
  // The LM Studio Python is appended at the END as a last-resort fallback only.

  if (IS_WIN) {
    const home = os.homedir();
    const localAppData = process.env.LOCALAPPDATA ?? path.join(home, "AppData", "Local");
    const programFiles = process.env.ProgramFiles ?? "C:\\Program Files";

    for (let minor = 15; minor >= minMinor; minor--) {
      const pad = minor.toString().padStart(2, "0");
      // python.org user installer (most common on Windows)
      candidates.push(path.join(localAppData, "Programs", "Python", `Python3${pad}`, "python.exe"));
      // python.org system installer
      candidates.push(path.join(programFiles, `Python3${pad}`, "python.exe"));
      candidates.push(`C:\\Python3${pad}\\python.exe`);
    }
    // Custom / pyenv-win paths
    candidates.push(path.join(localAppData, "Python", "bin", "python3.exe"));
    candidates.push(path.join(localAppData, "Python", "bin", "python.exe"));
    candidates.push(path.join(home, ".pyenv", "pyenv-win", "shims", "python.exe"));
  } else {
    for (let minor = 15; minor >= minMinor; minor--) {
      const ver = `${minMajor}.${minor}`;
      candidates.push(`/usr/local/bin/python${ver}`);                              // python.org
      candidates.push(`/opt/homebrew/bin/python${ver}`);                          // Homebrew (Apple Silicon)
      candidates.push(`/usr/local/opt/python@${ver}/bin/python${ver}`);           // Homebrew (Intel)
    }
    candidates.push("/Library/Developer/CommandLineTools/usr/bin/python3");        // Xcode CLT
    candidates.push("/usr/bin/python3");                                            // macOS / Linux system
  }

  // Last resort: LM Studio's own bundled Python. On macOS this Python has hardened
  // runtime and cannot load third-party native extensions with a different Team ID,
  // so we only fall back to it when nothing else is available (e.g. bare Windows
  // machines before a python.org installer has run).
  candidates.push(...lmsAmphibianPythonCandidates());

  return candidates;
}

function pythonActualVersion(pythonPath: string): [number, number] | null {
  try {
    const out = execFileSync(
      pythonPath,
      ["-c", "import sys; print(sys.version_info.major, sys.version_info.minor)"],
      { encoding: "utf-8", timeout: 5_000 }
    );
    const m = (out || "").trim().match(/^(\d+)\s+(\d+)$/);
    if (!m) return null;
    return [parseInt(m[1], 10), parseInt(m[2], 10)];
  } catch {
    return null;
  }
}

function versionMeetsMin(
  ver: [number, number],
  minMajor: number,
  minMinor: number
): boolean {
  return ver[0] > minMajor || (ver[0] === minMajor && ver[1] >= minMinor);
}

function pythonLabel(pythonPath: string): string {
  const p = pythonPath.replace(/\\/g, "/");
  if (p.includes("_amphibian"))  return "LM Studio bundled Python";
  if (p.includes("homebrew") || p.includes("/opt/homebrew") || p.includes("/usr/local/opt")) return "Homebrew Python";
  if (p.includes("/usr/local/bin"))    return "python.org Python (system-wide)";
  if (p.includes("Programs/Python") || p.includes("Program Files")) return "python.org Python (Windows installer)";
  if (p.includes("AppData"))           return "Python (user install)";
  if (p.includes("pyenv"))             return "pyenv Python";
  if (p.includes("CommandLineTools"))  return "Xcode CLT Python";
  if (p.includes("/usr/bin") || p === "python3" || p === "python") return "system Python";
  return "Python (unknown source)";
}

function findSystemPython(minMajor: number, minMinor: number): { path: string; label: string } {
  for (const candidate of systemPythonCandidates(minMajor, minMinor)) {
    if (!fs.existsSync(candidate)) continue;
    const ver = pythonActualVersion(candidate);
    if (ver && versionMeetsMin(ver, minMajor, minMinor)) {
      return { path: candidate, label: pythonLabel(candidate) };
    }
  }

  // Last-resort: try bare executable names via PATH (catches user-managed PATH entries
  // that LM Studio does not inherit, e.g. user-scoped python.org installs on Windows).
  for (const name of IS_WIN
    ? ["python3.exe", "python.exe", "python3", "python"]
    : ["python3", "python"]
  ) {
    const ver = pythonActualVersion(name);
    if (ver && versionMeetsMin(ver, minMajor, minMinor)) {
      return { path: name, label: pythonLabel(name) };
    }
  }

  const installHint = IS_WIN
    ? `Install Python ${minMajor}.${minMinor}+ from https://www.python.org/downloads/ ` +
      `(standard installer — not the Microsoft Store version). ` +
      `When the installer asks whether to add Python to PATH, answer Yes ` +
      `(this will not help LM Studio find Python directly, but causes no harm). ` +
      `Then restart LM Studio.`
    : `Install Python from https://www.python.org/downloads/ and restart LM Studio.`;
  throw new Error(
    `Python ${minMajor}.${minMinor}+ is required for the RAG environment but could not be found.\n` +
    installHint
  );
}

// ── Public API ─────────────────────────────────────────────────────────────────

/**
 * Ensure the .rag-venv environment exists and dependencies are installed.
 *
 * Fast path: if the ready marker exists AND the venv python is present,
 * returns immediately (no subprocess needed).
 *
 * Slow path (first run): creates the venv and installs python/requirements.txt.
 * Progress is streamed through onStatus — forward to ctx.status() in tools.
 */
export async function ensureRagVenv(
  onStatus: (msg: string) => void
): Promise<void> {
  // Ensure logs dir exists for the log file
  const logDir = logsDir();
  if (!fs.existsSync(logDir)) {
    fs.mkdirSync(logDir, { recursive: true });
  }
  const logFile = path.join(logDir, "rag-venv-setup.log");
  const ragImports = ["docling", "fitz"];

  // Always log which Python is in the venv (or will be used), for diagnostics.
  const logPythonInfo = () => {
    const venvPy = ragVenvPython();
    const version = ragVenvPythonVersion();
    // Resolve the symlink to find the real Python behind the venv
    let realPy = venvPy;
    try { realPy = fs.realpathSync(venvPy); } catch {}
    const label = pythonLabel(realPy);
    const line = `[python-in-venv] ${label}: ${realPy} (version: ${version ?? "unknown"})\n`;
    try { fs.appendFileSync(logFile, line); } catch {}
  };

  const ragReady = () =>
    fs.existsSync(ragReadyMarker()) &&
    ragVenvExists() &&
    fs.existsSync(ragVenvPip()) &&
    pythonImportsWorkLogged(ragVenvPython(), ragImports, logFile);

  // Fast path: marker present + python, pip and required imports present.
  if (ragVenvExists()) logPythonInfo();
  if (ragReady()) return;

  // Stale marker or broken env → remove and redo.
  if (fs.existsSync(ragReadyMarker())) {
    onStatus("RAG environment appears incomplete — rebuilding…");
    try { fs.rmSync(ragVenvDir(), { recursive: true, force: true }); } catch {}
    unlinkIfExists(ragReadyMarker());
  }

  // Outdated python in existing venv (< 3.9) → nuke and recreate
  if (ragVenvExists()) {
    const version = ragVenvPythonVersion();
    const parsed = version?.match(/^(\d+)\.(\d+)/);
    if (parsed) {
      const ver: [number, number] = [parseInt(parsed[1], 10), parseInt(parsed[2], 10)];
      if (!versionMeetsMin(ver, 3, 9)) {
        onStatus(`RAG environment uses Python ${version} — needs 3.9+, rebuilding…`);
        try { fs.rmSync(ragVenvDir(), { recursive: true, force: true }); } catch {}
        unlinkIfExists(ragReadyMarker());
      }
    }
  }

  if (ragVenvExists() && !fs.existsSync(ragVenvPip())) {
    onStatus("RAG environment appears incomplete — rebuilding…");
    try { fs.rmSync(ragVenvDir(), { recursive: true, force: true }); } catch {}
    unlinkIfExists(ragReadyMarker());
  }

  // Create venv if it still doesn't exist
  if (!ragVenvExists()) {
    onStatus("Setting up RAG Python environment (first run — this takes a moment)…");
    const { path: systemPython, label } = findSystemPython(3, 9);
    const msg = `Using ${label}: ${systemPython}`;
    onStatus(msg);
    try { fs.appendFileSync(logFile, `[python-selected] ${msg}\n`); } catch {}
    await runAndStream(
      systemPython,
      ["-m", "venv", ragVenvDir()],
      pluginRoot(),
      onStatus,
      logFile
    );
  }

  // Install / upgrade dependencies
  onStatus("Installing RAG dependencies (docling / PyMuPDF)…");
  // Use `python -m pip` instead of the pip executable to upgrade pip itself.
  // On Windows, pip.exe cannot replace itself while running (exit code 1).
  await runAndStream(
    ragVenvPython(),
    ["-m", "pip", "install", "--upgrade", "pip"],
    pluginRoot(),
    onStatus,
    logFile
  );
  await runAndStream(
    ragVenvPip(),
    ["install", "-r", requirementsFile()],
    pluginRoot(),
    (line) => {
      // Filter pip noise; only forward lines that look meaningful
      if (
        /Successfully installed|Collecting|Installing|already satisfied|Downloading/i.test(line)
      ) {
        onStatus(line);
      }
    },
    logFile
  );

  if (!pythonImportsWorkLogged(ragVenvPython(), ragImports, logFile)) {
    unlinkIfExists(ragReadyMarker());
    throw new Error("RAG environment setup completed, but required packages could not be imported. See rag-venv-setup.log for details.");
  }

  // Write ready marker
  fs.writeFileSync(ragReadyMarker(), new Date().toISOString() + "\n");
  onStatus("RAG Python environment ready.");
}