lmagent

import { tool, text } from "@lmstudio/sdk";
import { z } from "zod";
import { runProcess } from "./runProcess";
import { getPython } from "./pythonEnv";

function getDocsUrl(version: string) {
  const [major, minor] = version.split(".");
  return `https://docs.python.org/${major}.${minor}/`;
}

async function checkFreeThreading(py: string, cwd: string, timeout: number) {
  const code = `
import sys
print(getattr(sys, "_is_gil_enabled", "unknown"))
`;

  const result = await runProcess(py, ["-c", code], timeout, cwd);
  const out = result.stdout.trim();

  if (out === "False") return true;
  if (out === "True") return false;
  return null;
}

async function getBuildInfo(py: string, cwd: string, timeout: number) {
  const result = await runProcess(py, ["-c", "import sys; print(sys.version)"], timeout, cwd);
  return result.stdout.trim();
}

async function getPipVersion(py: string, cwd: string, timeout: number) {
  try {
    const result = await runProcess(py, ["-m", "pip", "--version"], timeout, cwd);
    return result.stdout.trim();
  } catch {
    return null;
  }
}

async function getPipList(py: string, cwd: string, timeout: number) {
  try {
    const result = await runProcess(py, ["-m", "pip", "list", "--format=json"], timeout, cwd);
    return JSON.parse(result.stdout);
  } catch {
    return [];
  }
}

function getFeatureHints(version: string) {
  const [major, minor] = version.split(".").map(Number);

  if (major === 3) {
    switch (minor) {
      case 14:
        return [
          "compression.zstd module",
          "asyncio introspection improvements",
          "t-string (Template) type",
        ];
      case 13:
        return [
          "experimental free-threading support",
          "legacy stdlib removals",
        ];
      case 12:
        return [
          "distutils removed",
          "f-string improvements",
          "typing enhancements",
        ];
      case 11:
        return [
          "exception notes",
          "exception groups",
          "match-case",
        ];
    }
  }

  return ["standard Python features"];
}

export function createPythonVersionTool(ctl: any) {
  return tool({
    name: "python_ver",
    description: text`
      Get Python runtime info (version, docs, features, pip, environment).
    `,
    parameters: {
      timeout_seconds: z.number().optional(),
    },

    implementation: async ({ timeout_seconds }) => {
      const cwd = ctl.getWorkingDirectory();
      const timeout = (timeout_seconds ?? 5) * 1000;
      
      const py = getPython();

      const result = await runProcess(py, ["--version"], timeout, cwd);
      const raw = result.stdout || result.stderr;

      const match = raw.match(/(\d+\.\d+\.\d+)/);
      if (!match) throw new Error("Failed to parse Python version");

      const version = match[1];

      const [freeThreading, build, pipVersion, pipList] = await Promise.all([
        checkFreeThreading(py, cwd, timeout),
        getBuildInfo(py, cwd, timeout),
        getPipVersion(py, cwd, timeout),
        getPipList(py, cwd, timeout),
      ]);

      return {
        version,
        docs_url: getDocsUrl(version),
        features: getFeatureHints(version),

        free_threading: freeThreading,
        build,
        pip_version: pipVersion,
        pip_list: pipList,

        python_executable: py,
      };
    },
  });
}