chonkie
Version:
🦛 CHONK your texts in TS with Chonkie!✨The no-nonsense lightweight and efficient chunking library.
105 lines • 5.03 kB
JavaScript
;
/** ChromaHandshake to integrate Chonkie with Chroma. */
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.ChromaHandshake = void 0;
const base_1 = require("./base");
const chromadb_1 = require("chromadb");
const utils_1 = require("./utils");
const base_2 = require("../types/base");
const uuid_1 = require("uuid");
/**
* ChromaHandshake to integrate Chonkie with Chroma.
*
* @param client - The ChromaClient to use.
* @param collectionName - The name of the collection to use.
* @param path - The path to the Chroma database. Can point to the running instance, Docker or Cloud.
* @param logLevel - The log level ('verbose' or 'silent'). Default: 'verbose'.
*/
class ChromaHandshake extends base_1.BaseHandshake {
constructor(client, collectionName, path, logLevel = 'verbose') {
super();
// If the client is not provided, create a new one
this.client = client !== null && client !== void 0 ? client : new chromadb_1.ChromaClient({ path });
// If the collection name is not provided, generate a random one
this.collectionName = collectionName !== null && collectionName !== void 0 ? collectionName : (0, utils_1.generateRandomCollectionName)();
this.logLevel = logLevel;
// Print to console the collection name if verbose
if (this.logLevel === 'verbose') {
console.log(`Using collection ${this.collectionName}`);
}
}
_getId(index, chunk) {
const id = (0, uuid_1.v5)(`CHUNK-${index}:${chunk.text}`, uuid_1.v5.DNS);
return id;
}
/**
* Write chunks to the collection provided in the constructor.
* @param chunks - The chunks to write.
*/
write(chunks) {
return __awaiter(this, void 0, void 0, function* () {
// Check if the collection exists and if not, create it
const collection = yield this.client.getOrCreateCollection({ name: this.collectionName });
// Create a list of ids and documents to upsert
const ids = [];
const documents = [];
const metadatas = [];
for (const [index, chunk] of chunks.entries()) {
ids.push(this._getId(index, chunk));
documents.push(chunk.text);
metadatas.push({
"start_index": chunk.startIndex,
"end_index": chunk.endIndex,
"token_count": chunk.tokenCount,
});
}
// Upsert the chunks into the collection
yield collection.upsert({
ids: ids,
documents: documents,
metadatas: metadatas,
});
// Print to console the number of chunks upserted if verbose
if (this.logLevel === 'verbose') {
console.log(`Upserted ${chunks.length} chunks into the collection ${this.collectionName}`);
}
});
}
/**
* Query the collection provided in the constructor.
* @param query - The query to search for.
* @param nResults - The number of results to return.
* @returns The chunks that match the query.
*/
query(query_1) {
return __awaiter(this, arguments, void 0, function* (query, nResults = 10) {
const collection = yield this.client.getCollection({ name: this.collectionName });
const results = yield collection.query({
queryTexts: [query],
nResults: nResults,
});
// Return the chunks
const { documents, metadatas } = results;
return documents[0].map((document, index) => {
const metadata = metadatas[0][index];
return new base_2.Chunk({
text: document !== null && document !== void 0 ? document : '',
startIndex: Number(metadata === null || metadata === void 0 ? void 0 : metadata.start_index) || 0,
endIndex: Number(metadata === null || metadata === void 0 ? void 0 : metadata.end_index) || 0,
tokenCount: Number(metadata === null || metadata === void 0 ? void 0 : metadata.token_count) || 0,
});
});
});
}
}
exports.ChromaHandshake = ChromaHandshake;
//# sourceMappingURL=chroma.js.map