Project Files
dist / index.js
"use strict";
/**
* @file index.ts
* @description Data-sampler LM Studio plugin entry point.
*
* Registers two tools via toolsProvider:
*
* Tool: sample
* Samples input text into optimised chunks using priority, relevance,
* or length-based strategies.
* See ./sampler.ts and ./strategies/ for implementation details.
*
* Tool: load_dataset
* Loads sample rows from HuggingFace datasets via the HF Datasets Server
* REST API. Supports config/split discovery, BM25 search, and multiple
* output formats (text, json, prompt).
* See ./dataset.ts for implementation details.
*
* @package data-sampler
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.main = main;
const sdk_1 = require("@lmstudio/sdk");
const zod_1 = require("zod");
const dataset_1 = require("./dataset");
const sampler_1 = require("./sampler");
// ---------------------------------------------------------------------------
// Plugin entry point
// ---------------------------------------------------------------------------
async function main(context) {
context.withToolsProvider(toolsProvider);
}
// ---------------------------------------------------------------------------
// Tools provider
// ---------------------------------------------------------------------------
const SAMPLE_TOOL_DESCRIPTION = "Sample input text into optimised chunks using one of three strategies.\n\n" +
"Strategies:\n" +
' - "priority": Scores chunks by auto-extracted keyword frequency and ' +
"position. Best when you want the most information-dense segments.\n" +
' - "relevance": BM25-style scoring against a provided query. ' +
"Best when you know what you are looking for.\n" +
' - "length": Simple fixed-size chunking with optional overlap. ' +
"Best when you need uniform segments.\n\n" +
"Use this tool to prepare large texts for downstream processing where " +
"context window limits matter.";
const LOAD_DATASET_TOOL_DESCRIPTION = "Load sample rows from a HuggingFace dataset via the " +
"HF Datasets Server API.\n\n" +
"Parameters:\n" +
" - dataset (required): HF dataset name, e.g. squad, imdb, " +
"ibm/duorc.\n" +
' - split (optional, default "train"): Which split to use.\n' +
" - config (optional): Dataset config/subset. Required when a " +
"dataset has multiple configs.\n" +
" - max_samples (optional, default 5, max 100): Rows to return.\n" +
" - query (optional): BM25 full-text search query.\n" +
' - format (optional, default "text"): Output format text, json, ' +
"or prompt.\n" +
" - template (optional): Custom template with {text} placeholder.\n\n" +
"Use this tool when you need sample data from HF datasets for " +
"few-shot examples, evaluation, or analysis without downloading " +
"the full dataset.";
async function toolsProvider(ctl) {
// -----------------------------------------------------------------------
// sample tool
// -----------------------------------------------------------------------
const sampleTool = (0, sdk_1.tool)({
name: "sample",
description: SAMPLE_TOOL_DESCRIPTION,
parameters: {
text: zod_1.z
.string()
.min(1, "Text must not be empty.")
.describe("The input text to sample into chunks."),
strategy: zod_1.z
.enum(["priority", "relevance", "length"])
.optional()
.default("priority")
.describe("Sampling strategy. 'priority' (default) scores by keyword " +
"frequency and position. 'relevance' requires a query. " +
"'length' produces fixed-size chunks."),
chunkSize: zod_1.z
.number()
.int()
.positive()
.optional()
.describe("Target character size per chunk. Default: 2000. Minimum: 1."),
overlap: zod_1.z
.number()
.int()
.min(0)
.optional()
.describe("Number of overlapping characters between consecutive chunks. " +
"Default: 0. Only applies when strategy is 'length'."),
maxChunks: zod_1.z
.number()
.int()
.positive()
.optional()
.describe("Maximum number of chunks to return. " +
"Default: unlimited (all chunks returned)."),
query: zod_1.z
.string()
.optional()
.describe("Query string for the 'relevance' strategy (BM25-style scoring). " +
"Ignored for other strategies."),
},
implementation: async ({ text, strategy, chunkSize, overlap, maxChunks, query }, { status, warn, signal }) => {
if (signal.aborted) {
return "Sampling cancelled by user.";
}
status(`Sampling ${text.length.toLocaleString()} chars using "${strategy}" strategy...`);
try {
const result = (0, sampler_1.sample)(text, {
strategy,
chunkSize,
overlap,
maxChunks,
query,
});
if (result.chunks.length === 0) {
return {
chunks: [],
totalChunks: 0,
originalLength: result.originalLength,
strategy: result.strategy,
message: "No chunks were produced. The input text may be empty.",
};
}
status(`Produced ${result.totalChunks} chunk(s) ` +
`from ${result.originalLength.toLocaleString()} chars.`);
return result;
}
catch (err) {
const message = err instanceof Error ? err.message : String(err ?? "unknown error");
warn(`Sampling error: ${message}`);
return {
chunks: [],
totalChunks: 0,
originalLength: text.length ?? 0,
strategy: strategy ?? "priority",
error: message,
};
}
},
});
// -----------------------------------------------------------------------
// load_dataset tool
// -----------------------------------------------------------------------
const loadDatasetTool = (0, sdk_1.tool)({
name: "load_dataset",
description: LOAD_DATASET_TOOL_DESCRIPTION,
parameters: {
dataset: zod_1.z
.string()
.min(1, "Dataset name must not be empty.")
.describe('HuggingFace dataset name, e.g. "squad", "imdb", ' +
'"ibm/duorc", or "bigcode/the-stack-dedup".'),
split: zod_1.z
.enum(["train", "test", "validation"])
.optional()
.default("train")
.describe('Dataset split to load from. Default: "train".'),
config: zod_1.z
.string()
.optional()
.describe("Dataset config/subset name. Required when a dataset has " +
"multiple configs (e.g. 'wikitext-103-raw-v1' for wiki text)."),
max_samples: zod_1.z
.number()
.int()
.min(1)
.max(100)
.optional()
.describe("Number of rows to return. Default: 5. Max: 100."),
query: zod_1.z
.string()
.optional()
.describe("Optional BM25 full-text search query to filter rows. " +
"When provided, the server returns the most relevant rows."),
format: zod_1.z
.enum(["text", "json", "prompt", "fewshot"])
.optional()
.default("text")
.describe('Output format. "text" (default): key:value lines. ' +
'"json": raw JSON objects. "prompt": === Sample N === format. ' +
'"fewshot": Input:/Output: pairs ready for prompt injection ' +
"(use input_columns/output_column to control mapping)."),
template: zod_1.z
.string()
.optional()
.describe("Custom template with {text} placeholder for each sample. " +
'Only used when format is "text".'),
indices: zod_1.z
.string()
.optional()
.describe('Comma/space-separated row indices to fetch, e.g. "0,5,10". ' +
"Takes precedence over max_samples. Only those rows are returned."),
input_columns: zod_1.z
.array(zod_1.z.string())
.optional()
.describe('Columns to treat as input ("fewshot" format only). ' +
"Defaults to all columns except output_column."),
output_column: zod_1.z
.string()
.optional()
.describe('Column to treat as output/answer ("fewshot" format only). ' +
"Defaults to the last column."),
},
implementation: async (params, { status, warn, signal }) => {
if (signal.aborted) {
return "Dataset loading cancelled by user.";
}
status(`Loading up to ${params.max_samples ?? 5} sample(s) ` +
`from dataset "${params.dataset}"...`);
try {
const result = await (0, dataset_1.loadDataset)(params, signal);
status(`Loaded ${result.count} sample(s) from ` +
`${result.dataset}/${result.config} ` +
`(split: ${result.split}, total rows: ${result.totalRows}).`);
return result;
}
catch (err) {
const message = err instanceof Error ? err.message : String(err ?? "unknown error");
warn(`Dataset loading error: ${message}`);
return {
dataset: params.dataset,
split: params.split ?? "train",
config: params.config ?? null,
samples: [],
totalRows: 0,
count: 0,
format: params.format ?? "text",
formatted: "",
error: message,
};
}
},
});
return [sampleTool, loadDatasetTool];
}
//# sourceMappingURL=index.js.map