langextract
Version:
A TypeScript library for extracting structured and grounded information from text using LLMs
250 lines • 10.3 kB
JavaScript
;
/**
* Copyright 2025 kmbro.
*
* This is a TypeScript translation of the original Python LangExtract library
* by Google LLC (https://github.com/google/langextract).
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.Annotator = exports.DocumentRepeatError = void 0;
/**
* Provides functionality for annotating text using a language model.
*/
const uuid_1 = require("uuid");
const types_1 = require("./types");
const tokenizer_1 = require("./tokenizer");
class DocumentRepeatError extends Error {
constructor(message) {
super(message);
this.name = "DocumentRepeatError";
}
}
exports.DocumentRepeatError = DocumentRepeatError;
const ATTRIBUTE_SUFFIX = "_attributes";
function mergeNonOverlappingExtractions(allExtractions) {
if (allExtractions.length === 0) {
return [];
}
if (allExtractions.length === 1) {
return allExtractions[0];
}
const mergedExtractions = [...allExtractions[0]];
for (let i = 1; i < allExtractions.length; i++) {
for (const extraction of allExtractions[i]) {
let overlaps = false;
if (extraction.charInterval) {
for (const existingExtraction of mergedExtractions) {
if (existingExtraction.charInterval) {
if (extractionsOverlap(extraction, existingExtraction)) {
overlaps = true;
break;
}
}
}
}
if (!overlaps) {
mergedExtractions.push(extraction);
}
}
}
return mergedExtractions;
}
function extractionsOverlap(extraction1, extraction2) {
if (!extraction1.charInterval || !extraction2.charInterval) {
return false;
}
const start1 = extraction1.charInterval.startPos ?? 0;
const end1 = extraction1.charInterval.endPos ?? 0;
const start2 = extraction2.charInterval.startPos ?? 0;
const end2 = extraction2.charInterval.endPos ?? 0;
return start1 < end2 && start2 < end1;
}
class Annotator {
constructor(languageModel, promptTemplate, options = {}) {
this.languageModel = languageModel;
this.promptTemplate = promptTemplate;
this.formatType = options.formatType ?? types_1.FormatType.YAML;
this.attributeSuffix = options.attributeSuffix ?? ATTRIBUTE_SUFFIX;
this.fenceOutput = options.fenceOutput ?? false;
this.maxTokens = options.maxTokens;
}
async annotateDocuments(documents, resolver, options = {}) {
const { maxCharBuffer = 200, batchLength = 1, debug = true, extractionPasses = 1 } = options;
if (extractionPasses === 1) {
return this.annotateDocumentsSinglePass(documents, resolver, {
maxCharBuffer,
batchLength,
debug,
});
}
else {
return this.annotateDocumentsSequentialPasses(documents, resolver, {
maxCharBuffer,
batchLength,
debug,
extractionPasses,
});
}
}
async annotateText(text, resolver, options = {}) {
const { maxCharBuffer = 200, batchLength = 1, additionalContext, debug = true, extractionPasses = 1 } = options;
const document = {
text,
additionalContext,
documentId: `doc_${(0, uuid_1.v4)().substring(0, 8)}`,
};
const documents = await this.annotateDocuments([document], resolver, {
maxCharBuffer,
batchLength,
debug,
extractionPasses,
});
return documents[0];
}
async annotateDocumentsSinglePass(documents, resolver, options) {
const { maxCharBuffer, batchLength, debug } = options;
const results = [];
// Process documents in batches
for (let i = 0; i < documents.length; i += batchLength) {
const batch = documents.slice(i, i + batchLength);
const batchResults = await this.processDocumentBatch(batch, resolver, {
maxCharBuffer,
debug,
});
results.push(...batchResults);
}
return results;
}
async annotateDocumentsSequentialPasses(documents, resolver, options) {
const { maxCharBuffer, batchLength, debug, extractionPasses } = options;
const results = [];
// Process documents in batches
for (let i = 0; i < documents.length; i += batchLength) {
const batch = documents.slice(i, i + batchLength);
const batchResults = await this.processDocumentBatchSequentialPasses(batch, resolver, {
maxCharBuffer,
debug,
extractionPasses,
});
results.push(...batchResults);
}
return results;
}
async processDocumentBatch(documents, resolver, options) {
const { maxCharBuffer } = options;
const results = [];
for (const document of documents) {
const chunks = this.chunkDocument(document, maxCharBuffer);
const allExtractions = [];
for (const chunk of chunks) {
const prompt = this.generatePrompt(chunk.text, document.additionalContext);
const modelOutputs = await this.languageModel.infer([prompt], {
maxDecodeSteps: this.maxTokens,
});
if (modelOutputs.length > 0 && modelOutputs[0].length > 0) {
const output = modelOutputs[0][0].output;
if (output) {
const extractions = resolver.resolve(output);
const alignedExtractions = resolver.align(extractions, chunk.text, chunk.tokenOffset, chunk.charOffset);
allExtractions.push(...alignedExtractions);
}
}
}
const annotatedDocument = {
documentId: document.documentId,
text: document.text,
extractions: allExtractions,
tokenizedText: document.tokenizedText ?? (0, tokenizer_1.tokenize)(document.text),
};
results.push(annotatedDocument);
}
return results;
}
async processDocumentBatchSequentialPasses(documents, resolver, options) {
const { maxCharBuffer, extractionPasses } = options;
const results = [];
for (const document of documents) {
const chunks = this.chunkDocument(document, maxCharBuffer);
const allPassExtractions = [];
// Perform multiple extraction passes
for (let pass = 0; pass < extractionPasses; pass++) {
const passExtractions = [];
for (const chunk of chunks) {
const prompt = this.generatePrompt(chunk.text, document.additionalContext);
const modelOutputs = await this.languageModel.infer([prompt], {
maxDecodeSteps: this.maxTokens,
});
if (modelOutputs.length > 0 && modelOutputs[0].length > 0) {
const output = modelOutputs[0][0].output;
if (output) {
const extractions = resolver.resolve(output);
const alignedExtractions = resolver.align(extractions, chunk.text, chunk.tokenOffset, chunk.charOffset);
passExtractions.push(...alignedExtractions);
}
}
}
allPassExtractions.push(passExtractions);
}
// Merge extractions from all passes
const mergedExtractions = mergeNonOverlappingExtractions(allPassExtractions);
const annotatedDocument = {
documentId: document.documentId,
text: document.text,
extractions: mergedExtractions,
tokenizedText: document.tokenizedText ?? (0, tokenizer_1.tokenize)(document.text),
};
results.push(annotatedDocument);
}
return results;
}
chunkDocument(document, maxCharBuffer) {
const text = document.text;
const chunks = [];
let currentPos = 0;
let tokenOffset = 0;
while (currentPos < text.length) {
const chunkEnd = Math.min(currentPos + maxCharBuffer, text.length);
const chunkText = text.substring(currentPos, chunkEnd);
chunks.push({
text: chunkText,
tokenOffset,
charOffset: currentPos,
});
currentPos = chunkEnd;
tokenOffset += (0, tokenizer_1.tokenize)(chunkText).tokens.length;
}
return chunks;
}
generatePrompt(text, additionalContext) {
// This is a simplified prompt generation
// In a real implementation, you'd use the QAPromptGenerator
const promptLines = [this.promptTemplate.description];
if (additionalContext) {
promptLines.push(additionalContext);
}
if (this.promptTemplate.examples.length > 0) {
promptLines.push("Examples:");
for (const example of this.promptTemplate.examples) {
promptLines.push(`Q: ${example.text}`);
promptLines.push(`A: ${JSON.stringify({ extractions: example.extractions })}`);
}
}
promptLines.push(`Q: ${text}`);
promptLines.push("A:");
return promptLines.join("\n");
}
}
exports.Annotator = Annotator;
//# sourceMappingURL=annotation.js.map