Project Files
src / utils / ragVenvSetup.ts
/**
* 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.");
}