diff --git a/package-lock.json b/package-lock.json index b271173..9f841aa 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10,9 +10,11 @@ "license": "Apache-2.0", "dependencies": { "@lmstudio/lms-common": "^0.3.0", + "@lmstudio/sdk": "^0.0.1", "boxen": "^5.1.2", "chalk": "^4.1.2", "cmd-ts": "^0.13.0", + "columnify": "^1.6.0", "inquirer": "^8.2.6", "zod": "^3.22.4" }, @@ -20,6 +22,7 @@ "lms-cli": "dist/index.js" }, "devDependencies": { + "@types/columnify": "^1.5.4", "@types/inquirer": "^9.0.7", "@types/node": "^20.12.5", "@typescript-eslint/eslint-plugin": "^6.20.0", @@ -186,6 +189,14 @@ "zod": "^3.22.4" } }, + "node_modules/@lmstudio/lms-isomorphic": { + "version": "0.1.0", + "resolved": "https://registry.npmjs.org/@lmstudio/lms-isomorphic/-/lms-isomorphic-0.1.0.tgz", + "integrity": "sha512-ofJNwIaYCLkFk1IkxKQNbZbwDU/s7RU8tnimu6BS7wt8vD5o/NXDaaYvEvDR6tdolVvblPvIKUxO7xZTONGDpw==", + "dependencies": { + "ws": "^8.16.0" + } + }, "node_modules/@lmstudio/lms-shared-types": { "version": "0.3.0", "resolved": "https://registry.npmjs.org/@lmstudio/lms-shared-types/-/lms-shared-types-0.3.0.tgz", @@ -194,6 +205,17 @@ "zod": "^3.22.4" } }, + "node_modules/@lmstudio/sdk": { + "version": "0.0.1", + "resolved": "https://registry.npmjs.org/@lmstudio/sdk/-/sdk-0.0.1.tgz", + "integrity": "sha512-dWaJdBdhE+z2V5G97nViZu4OQ6MdNal256EFc7RH5HQKEYGBG+Mcpvy5OPCgzS5AEv0yXKsgGGtM3RcTKrEP6w==", + "dependencies": { + "@lmstudio/lms-isomorphic": "^0.1.0", + "chalk": "^4.1.2", + "immer": "^10.0.4", + "zod": "^3.22.4" + } + }, "node_modules/@microsoft/tsdoc": { "version": "0.14.2", "resolved": "https://registry.npmjs.org/@microsoft/tsdoc/-/tsdoc-0.14.2.tgz", @@ -260,6 +282,12 @@ "node": ">= 8" } }, + "node_modules/@types/columnify": { + "version": "1.5.4", + "resolved": "https://registry.npmjs.org/@types/columnify/-/columnify-1.5.4.tgz", + "integrity": "sha512-YPEVzmy3kJupUee1ueLuvGspy6U2JHcxt6rYvRsSCEgVC54+KdBFjQ6NG/0koZk69e1bfXwSusgChwdFhvEXMw==", + "dev": true + }, "node_modules/@types/inquirer": { "version": "9.0.7", "resolved": "https://registry.npmjs.org/@types/inquirer/-/inquirer-9.0.7.tgz", @@ -983,6 +1011,18 @@ "resolved": "https://registry.npmjs.org/color-name/-/color-name-1.1.4.tgz", "integrity": "sha512-dOy+3AuW3a2wNbZHIuMZpTcgjGuLU/uBL/ubcZF9OXbDo8ff4O8yVp5Bf0efS8uEoYo5q4Fx7dY9OgQGXgAsQA==" }, + "node_modules/columnify": { + "version": "1.6.0", + "resolved": "https://registry.npmjs.org/columnify/-/columnify-1.6.0.tgz", + "integrity": "sha512-lomjuFZKfM6MSAnV9aCZC9sc0qGbmZdfygNv+nCpqVkSKdCxCklLtd16O0EILGkImHw9ZpHkAnHaB+8Zxq5W6Q==", + "dependencies": { + "strip-ansi": "^6.0.1", + "wcwidth": "^1.0.0" + }, + "engines": { + "node": ">=8.0.0" + } + }, "node_modules/concat-map": { "version": "0.0.1", "resolved": "https://registry.npmjs.org/concat-map/-/concat-map-0.0.1.tgz", @@ -3739,6 +3779,26 @@ "integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==", "dev": true }, + "node_modules/ws": { + "version": "8.16.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.16.0.tgz", + "integrity": "sha512-HS0c//TP7Ina87TfiPUz1rQzMhHrl/SG2guqRcTOIUYD2q8uhUdNHZYJUaQ8aTGPzCh+c6oawMKW35nFl1dxyQ==", + "engines": { + "node": ">=10.0.0" + }, + "peerDependencies": { + "bufferutil": "^4.0.1", + "utf-8-validate": ">=5.0.2" + }, + "peerDependenciesMeta": { + "bufferutil": { + "optional": true + }, + "utf-8-validate": { + "optional": true + } + } + }, "node_modules/yallist": { "version": "4.0.0", "resolved": "https://registry.npmjs.org/yallist/-/yallist-4.0.0.tgz", diff --git a/package.json b/package.json index f60c3da..76e1c7a 100644 --- a/package.json +++ b/package.json @@ -14,13 +14,16 @@ "license": "Apache-2.0", "dependencies": { "@lmstudio/lms-common": "^0.3.0", + "@lmstudio/sdk": "^0.0.1", "boxen": "^5.1.2", "chalk": "^4.1.2", "cmd-ts": "^0.13.0", + "columnify": "^1.6.0", "inquirer": "^8.2.6", "zod": "^3.22.4" }, "devDependencies": { + "@types/columnify": "^1.5.4", "@types/inquirer": "^9.0.7", "@types/node": "^20.12.5", "@typescript-eslint/eslint-plugin": "^6.20.0", diff --git a/src/createClient.ts b/src/createClient.ts new file mode 100644 index 0000000..ed8dec1 --- /dev/null +++ b/src/createClient.ts @@ -0,0 +1,8 @@ +import { type SimpleLogger } from "@lmstudio/lms-common"; +import { LMStudioClient } from "@lmstudio/sdk"; + +export function createClient(logger: SimpleLogger) { + return new LMStudioClient({ + logger, + }); +} diff --git a/src/formatSizeBytes1000.ts b/src/formatSizeBytes1000.ts new file mode 100644 index 0000000..58f2048 --- /dev/null +++ b/src/formatSizeBytes1000.ts @@ -0,0 +1,35 @@ +import chalk from "chalk"; + +const kB = 1000; +const mB = 1000 * kB; +const gB = 1000 * mB; +const tB = 1000 * gB; +const grayThreshold = 128 * mB; +const greenThreshold = 1 * gB; +const yellowThreshold = 24 * gB; + +export function formatSizeBytes1000(sizeBytes: number) { + if (sizeBytes < kB) { + return `${sizeBytes} B`; + } else if (sizeBytes < mB) { + return `${(sizeBytes / kB).toFixed(2)} KB`; + } else if (sizeBytes < gB) { + return `${(sizeBytes / mB).toFixed(2)} MB`; + } else if (sizeBytes < tB) { + return `${(sizeBytes / gB).toFixed(2)} GB`; + } else { + return `${(sizeBytes / tB).toFixed(2)} TB`; + } +} + +export function formatSizeBytesWithColor1000(sizeBytes: number) { + if (sizeBytes < grayThreshold) { + return chalk.gray(formatSizeBytes1000(sizeBytes)); + } else if (sizeBytes < greenThreshold) { + return chalk.green(formatSizeBytes1000(sizeBytes)); + } else if (sizeBytes < yellowThreshold) { + return chalk.yellow(formatSizeBytes1000(sizeBytes)); + } else { + return chalk.red(formatSizeBytes1000(sizeBytes)); + } +} diff --git a/src/index.ts b/src/index.ts index ab60262..6345ef6 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,4 +1,5 @@ import { run, subcommands } from "cmd-ts"; +import { list } from "./subcommands/list"; import { start, status, stop } from "./subcommands/server"; import { printVersion, version } from "./subcommands/version"; @@ -15,6 +16,7 @@ const cli = subcommands({ start, status, stop, + list, }, }); diff --git a/src/subcommands/list.ts b/src/subcommands/list.ts new file mode 100644 index 0000000..e8e9814 --- /dev/null +++ b/src/subcommands/list.ts @@ -0,0 +1,147 @@ +import { type DownloadedModel } from "@lmstudio/sdk"; +import chalk from "chalk"; +import { command, subcommands } from "cmd-ts"; +import columnify from "columnify"; +import { createClient } from "../createClient"; +import { formatSizeBytesWithColor1000 } from "../formatSizeBytes1000"; +import { createLogger, logLevelArgs } from "../logLevel"; + +function loadedCheck(count: number) { + if (count === 0) { + return ""; + } else if (count === 1) { + return chalk.bgGreenBright.black(" ✓ LOADED "); + } else { + return chalk.bgGreenBright.black(` ✓ LOADED (${count}) `); + } +} + +function coloredArch(arch?: string) { + return arch ?? ""; +} + +function printDownloadedModelsTable( + title: string, + downloadedModels: Array, + loadedModels: Array<{ address: string; identifier: string }>, +) { + interface DownloadedModelWithExtraInfo extends DownloadedModel { + loadedIdentifiers: Array; + group: string; + remaining: string; + } + const downloadedModelsGroups = downloadedModels + // Attach 1) all the loadedIdentifiers 2) group name (user/repo) to each model + .map(model => { + const segments = model.address.split("/"); + return { + ...model, + loadedIdentifiers: loadedModels + .filter(loadedModel => loadedModel.address === model.address) + .map(loadedModel => loadedModel.identifier), + group: segments.slice(0, 2).join("/"), + remaining: segments.slice(2).join("/"), + }; + }) + // Group by group name into a map + .reduce((acc, model) => { + let group = acc.get(model.group); + if (!group) { + group = []; + acc.set(model.group, group); + } + group.push(model); + return acc; + }, new Map>()); + + const downloadedModelsAndHeadlines = [...downloadedModelsGroups.entries()].flatMap( + ([group, models]) => { + if (models.length === 1 && models[0].remaining === "") { + // Group is a model itself + const model = models[0]; + return { + address: chalk.whiteBright(group), + sizeBytes: formatSizeBytesWithColor1000(model.sizeBytes), + arch: coloredArch(model.architecture), + loaded: loadedCheck(model.loadedIdentifiers.length), + }; + } + return [ + // Group title + { address: chalk.whiteBright(group), sizeBytes: "", arch: "", loaded: "" }, + // Models + ...models.map(model => ({ + address: chalk.black(". ") + chalk.gray("/" + model.remaining), + sizeBytes: formatSizeBytesWithColor1000(model.sizeBytes), + arch: coloredArch(model.architecture), + loaded: loadedCheck(model.loadedIdentifiers.length), + })), + ]; + }, + ); + + console.info(title); + console.info(); + console.info( + columnify(downloadedModelsAndHeadlines, { + columns: ["address", "sizeBytes", "arch", "loaded"], + config: { + address: { + headingTransform: () => chalk.cyanBright("ADDRESS"), + }, + sizeBytes: { + headingTransform: () => chalk.cyanBright("SIZE"), + align: "right", + }, + arch: { + headingTransform: () => chalk.cyanBright("ARCHITECTURE"), + align: "left", + }, + loaded: { + headingTransform: () => chalk.cyanBright("LOADED"), + align: "left", + }, + }, + preserveNewLines: true, + columnSplitter: " ", + }), + ); +} + +const downloaded = command({ + name: "downloaded", + description: "List downloaded models", + args: { + ...logLevelArgs, + }, + handler: async args => { + const logger = createLogger(args); + const client = createClient(logger); + + const downloadedModels = await client.system.listDownloadedModels(); + const loadedModels = await client.llm.listLoaded(); + + console.info(); + console.info(); + printDownloadedModelsTable( + chalk.bgGreenBright.black(" LLM ") + " " + chalk.green("(Large Language Models)"), + downloadedModels.filter(model => model.type === "llm"), + loadedModels, + ); + console.info(); + console.info(); + printDownloadedModelsTable( + chalk.bgGreenBright.black(" Embeddings "), + downloadedModels.filter(model => model.type === "embedding"), + loadedModels, + ); + console.info(); + console.info(); + }, +}); + +export const list = subcommands({ + name: "list", + description: "List models", + cmds: { downloaded }, +});