python-tools

import { randomUUID } from "node:crypto";
import { mkdir, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";

import {
  type PythonRunProcessResult,
  resolvePythonCommand,
  runResolvedPythonCommand,
} from "./pythonResolver";

export type PythonRunResult = PythonRunProcessResult;

const DEFAULT_TIMEOUT_SECONDS = 5;
const MAX_TIMEOUT_SECONDS = 20;
const MAX_CODE_LENGTH = 20_000;

export async function runPythonCode(
  code: string,
  timeoutSeconds = DEFAULT_TIMEOUT_SECONDS,
): Promise<PythonRunResult> {
  validatePythonRequest(code, timeoutSeconds);

  const tempDir = path.join(tmpdir(), "lmstudio-python-tools");
  const filePath = path.join(tempDir, `run-${randomUUID()}.py`);

  await mkdir(tempDir, { recursive: true });
  await writeFile(filePath, code, "utf8");

  try {
    return await runPythonFile(filePath, timeoutSeconds);
  } finally {
    await rm(filePath, { force: true });
  }
}

function validatePythonRequest(code: string, timeoutSeconds: number): void {
  if (code.trim().length === 0) {
    throw new Error("Python code is required.");
  }

  if (code.length > MAX_CODE_LENGTH) {
    throw new Error(`Python code must be ${MAX_CODE_LENGTH} characters or fewer.`);
  }

  if (!Number.isFinite(timeoutSeconds) || timeoutSeconds <= 0) {
    throw new Error("timeoutSeconds must be a positive number.");
  }

  if (timeoutSeconds > MAX_TIMEOUT_SECONDS) {
    throw new Error(`timeoutSeconds must be ${MAX_TIMEOUT_SECONDS} seconds or fewer.`);
  }
}

async function runPythonFile(filePath: string, timeoutSeconds: number): Promise<PythonRunResult> {
  const pythonCommand = await resolvePythonCommand();
  return await runResolvedPythonCommand(pythonCommand, [filePath], timeoutSeconds);
}