lmagent

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

export function createPipManagerTool(ctl: any) {
  return tool({
    name: "pip_manage",
    description: text`
      Manage Python packages using system Python (user-level install).
    `,
    parameters: {
      action: z.enum([
        "install",
        "uninstall",
        "list",
        "show",
        "freeze",
        "upgrade"
      ]),
      packages: z.array(z.string()).optional(),
      timeout_seconds: z.number().optional(),
    },

    implementation: async ({ action, packages, timeout_seconds }) => {
      const cwd = ctl.getWorkingDirectory();
      const timeout = (timeout_seconds ?? 60) * 1000;

      const python = getPython();

      let args: string[] = ["-m", "pip"];

      switch (action) {
        case "install":
          if (!packages?.length) throw new Error("No packages provided");
          args.push(
            "install",
            "--user",
            "--disable-pip-version-check",
            "--no-input",
            ...packages
          );
          break;

        case "uninstall":
          if (!packages?.length) throw new Error("No packages provided");
          args.push("uninstall", "-y", ...packages);
          break;

        case "list":
          args.push("list", "--format=json");
          break;

        case "show":
          if (!packages?.length) throw new Error("No packages provided");
          args.push("show", ...packages);
          break;

        case "freeze":
          args.push("freeze");
          break;

        case "upgrade":
          if (!packages?.length) throw new Error("No packages provided");
          args.push("install", "--upgrade", "--user", ...packages);
          break;

        default:
          throw new Error("Invalid action");
      }

      const result = await runProcess(python, args, timeout, cwd);

      if (action === "list") {
        try {
          return {
            packages: JSON.parse(result.stdout),
          };
        } catch {
          return { raw: result.stdout };
        }
      }

      return {
        stdout: result.stdout,
        stderr: result.stderr,
      };
    },
  });
}