chonkie
Version:
🦛 CHONK your texts in TS with Chonkie!✨The no-nonsense lightweight and efficient chunking library.
378 lines • 16.8 kB
JavaScript
;
/** Module containing CodeChunker class. */
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.CodeChunker = void 0;
const web_tree_sitter_1 = require("web-tree-sitter");
const fs = __importStar(require("fs"));
const path = __importStar(require("path"));
const tokenizer_1 = require("../tokenizer");
const code_1 = require("../types/code");
const base_1 = require("./base");
/**
* CodeChunker class extends BaseChunker and provides functionality for chunking code.
*/
class CodeChunker extends base_1.BaseChunker {
/**
* Private constructor. Use `CodeChunker.create()` to instantiate.
*/
constructor(tokenizer, chunkSize, lang, includeNodes = false) {
super(tokenizer);
this.parser = null;
this.language = undefined;
if (chunkSize <= 0) {
throw new Error("chunkSize must be greater than 0");
}
this.chunkSize = chunkSize;
this.lang = lang;
this.includeNodes = includeNodes;
}
/**
* Creates and initializes a CodeChunker instance that is directly callable.
*/
static create() {
return __awaiter(this, arguments, void 0, function* (options = {}) {
if (!CodeChunker.treeSitterInitialized && web_tree_sitter_1.Parser.init) {
try {
yield web_tree_sitter_1.Parser.init();
CodeChunker.treeSitterInitialized = true;
}
catch (error) {
// Log error but don't necessarily throw, as some environments might not need/support it
// or tree-sitter might still work for some languages.
console.error("Failed to run Parser.init():", error);
}
}
const { tokenizer = "Xenova/gpt2", chunkSize = 512, lang, includeNodes = false } = options;
let tokenizerInstance;
if (typeof tokenizer === 'string') {
tokenizerInstance = yield tokenizer_1.Tokenizer.create(tokenizer);
}
else {
tokenizerInstance = tokenizer;
}
if (!lang) {
throw new Error("Language must be specified for code chunking");
}
const plainInstance = new CodeChunker(tokenizerInstance, chunkSize, lang, includeNodes);
// Create the callable function wrapper
const callableFn = function (textOrTexts, showProgress) {
if (typeof textOrTexts === 'string') {
return plainInstance.call(textOrTexts, showProgress);
}
else {
return plainInstance.call(textOrTexts, showProgress);
}
};
// Set the prototype so that 'instanceof CodeChunker' works
Object.setPrototypeOf(callableFn, CodeChunker.prototype);
// Copy all enumerable own properties from plainInstance to callableFn
Object.assign(callableFn, plainInstance);
return callableFn;
});
}
/**
* Recursively finds the nearest node_modules directory from a starting directory.
* @param startDir The directory to start searching from.
* @returns The absolute path to the node_modules directory, or null if not found.
*/
static findNearestNodeModules(startDir) {
let dir = path.resolve(startDir); // Ensure absolute path
while (true) {
const candidate = path.join(dir, "node_modules");
if (fs.existsSync(candidate) && fs.statSync(candidate).isDirectory()) {
return candidate;
}
const parent = path.dirname(dir);
if (parent === dir)
break; // Reached filesystem root
dir = parent;
}
return null;
}
/**
* Initialize the tree-sitter parser for the given language using WASM.
*/
_initParser(lang) {
return __awaiter(this, void 0, void 0, function* () {
if (this.parser && this.language) {
return; // Already initialized for this instance
}
// Parser.init() is now called in the static create method
// and treeSitterInitialized is managed there.
// Convert language name to lowercase and replace hyphens with underscores
const formattedLang = lang.toLowerCase().replace(/-/g, '_');
// Find node_modules starting from the current file's directory and searching upwards
const nodeModulesPath = CodeChunker.findNearestNodeModules(__dirname);
if (!nodeModulesPath) {
throw new Error("node_modules directory not found. " +
"This is required for loading tree-sitter language WASM files from 'tree-sitter-wasms' package.");
}
const wasmPath = path.join(nodeModulesPath, `tree-sitter-wasms/out/tree-sitter-${formattedLang}.wasm`);
if (!fs.existsSync(wasmPath)) {
throw new Error(`Tree-sitter WASM file for language "${formattedLang}" not found at ${wasmPath}. ` +
`Ensure 'tree-sitter-wasms' package is installed and the language is supported.`);
}
try {
const wasmBuffer = fs.readFileSync(wasmPath);
this.language = yield web_tree_sitter_1.Language.load(wasmBuffer);
this.parser = new web_tree_sitter_1.Parser();
this.parser.setLanguage(this.language);
}
catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error);
throw new Error(`Failed to initialize tree-sitter parser for language "${lang}" from WASM path "${wasmPath}": ${errorMessage}`);
}
});
}
/**
* Merge node groups together.
*/
_mergeNodeGroups(nodeGroups) {
return nodeGroups.flat();
}
/**
* Group child nodes based on their token counts.
*/
_groupChildNodes(node) {
return __awaiter(this, void 0, void 0, function* () {
if (!node.children || node.children.length === 0) {
return [[], []];
}
const nodeGroups = [];
const groupTokenCounts = [];
let currentTokenCount = 0;
let currentNodeGroup = [];
for (const child of node.children) {
const childText = child.text;
const tokenCount = yield this.tokenizer.countTokens(childText);
if (tokenCount > this.chunkSize) {
if (currentNodeGroup.length > 0) {
nodeGroups.push(currentNodeGroup);
groupTokenCounts.push(currentTokenCount);
currentNodeGroup = [];
currentTokenCount = 0;
}
const [childGroups, childTokenCounts] = yield this._groupChildNodes(child);
nodeGroups.push(...childGroups);
groupTokenCounts.push(...childTokenCounts);
}
else if (currentTokenCount + tokenCount > this.chunkSize) {
nodeGroups.push(currentNodeGroup);
groupTokenCounts.push(currentTokenCount);
currentNodeGroup = [child];
currentTokenCount = tokenCount;
}
else {
currentNodeGroup.push(child);
currentTokenCount += tokenCount;
}
}
if (currentNodeGroup.length > 0) {
nodeGroups.push(currentNodeGroup);
groupTokenCounts.push(currentTokenCount);
}
// Calculate cumulative token counts for optimal grouping
const cumulativeTokenCounts = [0];
for (const count of groupTokenCounts) {
cumulativeTokenCounts.push(cumulativeTokenCounts[cumulativeTokenCounts.length - 1] + count);
}
// Merge groups optimally using binary search
const mergedNodeGroups = [];
const mergedTokenCounts = [];
let pos = 0;
while (pos < nodeGroups.length) {
const startCumulativeCount = cumulativeTokenCounts[pos];
const requiredCumulativeTarget = startCumulativeCount + this.chunkSize;
// Find the optimal split point using binary search
let index = this._bisectLeft(cumulativeTokenCounts, requiredCumulativeTarget, pos) - 1;
index = Math.min(index, nodeGroups.length);
// Handle edge cases
if (index === pos) {
index = pos + 1;
}
// Merge the groups
const groupsToMerge = nodeGroups.slice(pos, index);
mergedNodeGroups.push(this._mergeNodeGroups(groupsToMerge));
// Calculate the actual token count for this merged group
const actualMergedCount = cumulativeTokenCounts[index] - cumulativeTokenCounts[pos];
mergedTokenCounts.push(actualMergedCount);
pos = index;
}
return [mergedNodeGroups, mergedTokenCounts];
});
}
/**
* Binary search to find the first index where the value is greater than or equal to the target.
*/
_bisectLeft(arr, target, lo = 0) {
let hi = arr.length;
while (lo < hi) {
const mid = Math.floor((lo + hi) / 2);
if (arr[mid] < target) {
lo = mid + 1;
}
else {
hi = mid;
}
}
return lo;
}
/**
* Get texts from node groups using original byte offsets.
*/
_getTextsFromNodeGroups(nodeGroups, originalTextBytes) {
var _a, _b, _c;
const chunkTexts = [];
for (let i = 0; i < nodeGroups.length; i++) {
const group = nodeGroups[i];
if (!group.length)
continue;
const startNode = group[0];
const endNode = group[group.length - 1];
let startByte = startNode.startIndex;
let endByte = endNode.endIndex;
if (startByte > endByte) {
console.warn(`Warning: Skipping group due to invalid byte order. Start: ${startByte}, End: ${endByte}`);
continue;
}
if (startByte < 0 || endByte > originalTextBytes.length) {
console.warn(`Warning: Skipping group due to out-of-bounds byte offsets. Start: ${startByte}, End: ${endByte}, Text Length: ${originalTextBytes.length}`);
continue;
}
if (i < nodeGroups.length - 1) {
endByte = nodeGroups[i + 1][0].startIndex;
}
try {
const chunkBytes = originalTextBytes.slice(startByte, endByte);
const text = chunkBytes.toString('utf-8');
chunkTexts.push(text);
}
catch (error) {
console.warn(`Warning: Error decoding bytes for chunk (${startByte}-${endByte}): ${error}`);
chunkTexts.push("");
}
}
// Add any missing bytes at the start and end
if (((_b = (_a = nodeGroups[0]) === null || _a === void 0 ? void 0 : _a[0]) === null || _b === void 0 ? void 0 : _b.startIndex) > 0) {
const initialBytes = originalTextBytes.slice(0, nodeGroups[0][0].startIndex);
chunkTexts[0] = initialBytes.toString('utf-8') + chunkTexts[0];
}
const lastGroup = nodeGroups[nodeGroups.length - 1];
if (((_c = lastGroup === null || lastGroup === void 0 ? void 0 : lastGroup[lastGroup.length - 1]) === null || _c === void 0 ? void 0 : _c.endIndex) < originalTextBytes.length) {
const remainingBytes = originalTextBytes.slice(lastGroup[lastGroup.length - 1].endIndex);
chunkTexts[chunkTexts.length - 1] += remainingBytes.toString('utf-8');
}
return chunkTexts;
}
/**
* Create CodeChunk objects from texts, token counts, and node groups.
*/
_createChunks(texts, tokenCounts, nodeGroups) {
const chunks = [];
let currentIndex = 0;
for (let i = 0; i < texts.length; i++) {
const text = texts[i];
const tokenCount = tokenCounts[i];
const nodeGroup = this.includeNodes ? nodeGroups[i] : undefined;
chunks.push(new code_1.CodeChunk({
text,
startIndex: currentIndex,
endIndex: currentIndex + text.length,
tokenCount,
lang: this.lang,
nodes: nodeGroup
}));
currentIndex += text.length;
}
return chunks;
}
/**
* Recursively chunks the code based on context from tree-sitter.
*/
chunk(text) {
return __awaiter(this, void 0, void 0, function* () {
if (!text.trim()) {
return [];
}
const originalTextBytes = Buffer.from(text, 'utf-8');
if (!this.lang) {
throw new Error("Language must be specified for code chunking");
}
yield this._initParser(this.lang);
if (!this.parser) {
throw new Error("Parser not initialized");
}
let tree = null;
try {
tree = this.parser.parse(originalTextBytes.toString());
if (!tree) {
throw new Error("Failed to parse text");
}
const rootNode = tree.rootNode;
const [nodeGroups, tokenCounts] = yield this._groupChildNodes(rootNode);
const texts = this._getTextsFromNodeGroups(nodeGroups, originalTextBytes);
return this._createChunks(texts, tokenCounts, nodeGroups);
}
finally {
// No need to explicitly delete the tree - it will be garbage collected
if (!this.includeNodes) {
tree = null;
}
}
});
}
/**
* Return a string representation of the CodeChunker.
*/
toString() {
return `CodeChunker(tokenizer=${this.tokenizer}, ` +
`chunkSize=${this.chunkSize}, ` +
`lang=${this.lang}, ` +
`includeNodes=${this.includeNodes})`;
}
}
exports.CodeChunker = CodeChunker;
CodeChunker.treeSitterInitialized = false;
//# sourceMappingURL=code.js.map