langextract
Version:
A TypeScript library for extracting structured and grounded information from text using LLMs
317 lines • 11.6 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.
*/
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.OllamaLanguageModel = exports.OpenAILanguageModel = exports.GeminiLanguageModel = exports.InferenceOutputError = void 0;
/**
* Simple library for performing language model inference.
*/
const axios_1 = __importDefault(require("axios"));
const types_1 = require("./types");
class InferenceOutputError extends Error {
constructor(message) {
super(message);
this.name = "InferenceOutputError";
}
}
exports.InferenceOutputError = InferenceOutputError;
class GeminiLanguageModel {
constructor(config = {}) {
this.config = {
modelId: "gemini-2.5-flash",
apiKey: "",
formatType: types_1.FormatType.JSON,
temperature: 0.0,
maxWorkers: 10,
maxTokens: 2048,
...config,
};
this.constraint = { constraintType: "none" };
}
async infer(batchPrompts, options = {}) {
const results = [];
for (const prompt of batchPrompts) {
try {
const result = await this.processSinglePrompt(prompt, options);
results.push([result]);
}
catch (error) {
console.error("Error processing prompt:", error);
results.push([{ score: 0, output: undefined }]);
}
}
return results;
}
async processSinglePrompt(prompt, options) {
const config = {
temperature: options.temperature ?? this.config.temperature,
maxOutputTokens: options.maxDecodeSteps ?? this.config.maxTokens ?? 2048,
...options,
};
try {
const response = await this.callGeminiAPI(prompt, config);
return {
score: 1.0,
output: response,
};
}
catch (error) {
throw new InferenceOutputError(`Failed to get response from Gemini: ${error}`);
}
}
async callGeminiAPI(prompt, config) {
const baseUrl = this.config.modelUrl || "https://generativelanguage.googleapis.com";
const url = `${baseUrl}/v1beta/models/${this.config.modelId}:generateContent`;
const requestBody = {
contents: [
{
parts: [
{
text: prompt,
},
],
},
],
generationConfig: {
temperature: config.temperature,
maxOutputTokens: config.maxOutputTokens,
},
};
// Add schema if available
if (this.config.geminiSchema) {
requestBody.generationConfig.responseSchema = this.config.geminiSchema.schemaDict;
}
try {
const response = await axios_1.default.post(url, requestBody, {
headers: {
"Content-Type": "application/json",
"x-goog-api-key": this.config.apiKey,
},
timeout: 30000,
});
if (response.data.candidates && response.data.candidates[0]?.content?.parts?.[0]?.text) {
return response.data.candidates[0].content.parts[0].text;
}
else {
throw new Error("Invalid response format from Gemini API");
}
}
catch (error) {
if (axios_1.default.isAxiosError(error)) {
throw new Error(`Gemini API error: ${error.response?.data?.error?.message || error.message}`);
}
throw error;
}
}
parseOutput(output) {
try {
if (this.config.formatType === types_1.FormatType.JSON) {
return JSON.parse(output);
}
else {
// For YAML, you would need a YAML parser
// For now, return the raw output
return output;
}
}
catch (error) {
throw new Error(`Failed to parse output: ${error}`);
}
}
}
exports.GeminiLanguageModel = GeminiLanguageModel;
class OpenAILanguageModel {
constructor(config = {}) {
this.config = {
model: "gpt-4o-mini",
apiKey: "",
formatType: types_1.FormatType.JSON,
temperature: 0.0,
maxWorkers: 10,
baseURL: "https://api.openai.com/v1",
maxTokens: 2048,
...config,
};
this.constraint = { constraintType: "none" };
}
async infer(batchPrompts, options = {}) {
const results = [];
for (const prompt of batchPrompts) {
try {
const result = await this.processSinglePrompt(prompt, options);
results.push([result]);
}
catch (error) {
console.error("Error processing prompt:", error);
results.push([{ score: 0, output: undefined }]);
}
}
return results;
}
async processSinglePrompt(prompt, options) {
const config = {
temperature: options.temperature ?? this.config.temperature,
maxTokens: options.maxDecodeSteps ?? this.config.maxTokens ?? 2048,
...options,
};
try {
const response = await this.callOpenAIAPI(prompt, config);
return {
score: 1.0,
output: response,
};
}
catch (error) {
throw new InferenceOutputError(`Failed to get response from OpenAI: ${error}`);
}
}
async callOpenAIAPI(prompt, config) {
const url = `${this.config.baseURL}/chat/completions`;
const requestBody = {
model: this.config.model,
messages: [
{
role: "user",
content: prompt + "\n\n" + "Return the response in JSON format.",
},
],
temperature: config.temperature,
max_tokens: config.maxTokens,
};
// Add JSON response format when formatType is JSON
// OpenAI requires the word "json" in the prompt when using response_format: { type: "json_object" }
if (this.config.formatType === types_1.FormatType.JSON) {
requestBody.response_format = { type: "json_object" };
}
// Add function calling for schema enforcement if schema is available
if (this.config.openAISchema) {
requestBody.tools = [
{
type: "function",
function: {
name: "extract_data",
description: "Extract structured data from the text according to the schema",
parameters: this.config.openAISchema.schemaDict,
},
},
];
requestBody.tool_choice = {
type: "function",
function: { name: "extract_data" },
};
}
try {
const response = await axios_1.default.post(url, requestBody, {
headers: {
"Content-Type": "application/json",
Authorization: `Bearer ${this.config.apiKey}`,
},
timeout: 30000,
});
if (response.data.choices && response.data.choices[0]?.message?.content) {
return response.data.choices[0].message.content;
}
else if (response.data.choices && response.data.choices[0]?.message?.tool_calls) {
// Handle function call response
const toolCall = response.data.choices[0].message.tool_calls[0];
if (toolCall && toolCall.function && toolCall.function.arguments) {
return toolCall.function.arguments;
}
}
throw new Error("Invalid response format from OpenAI API");
}
catch (error) {
if (axios_1.default.isAxiosError(error)) {
throw new Error(`OpenAI API error: ${error.response?.data?.error?.message || error.message}`);
}
throw error;
}
}
parseOutput(output) {
try {
if (this.config.formatType === types_1.FormatType.JSON) {
return JSON.parse(output);
}
else {
// For YAML, you would need a YAML parser
// For now, return the raw output
return output;
}
}
catch (error) {
throw new Error(`Failed to parse output: ${error}`);
}
}
}
exports.OpenAILanguageModel = OpenAILanguageModel;
class OllamaLanguageModel {
constructor(config = {}) {
this.config = {
model: "gemma2:latest",
modelUrl: "http://localhost:11434",
structuredOutputFormat: "json",
temperature: 0.8,
maxTokens: 2048,
...config,
};
this.constraint = { constraintType: "none" };
}
async infer(batchPrompts, options = {}) {
const results = [];
for (const prompt of batchPrompts) {
try {
const result = await this.ollamaQuery(prompt, options);
results.push([{ score: 1.0, output: result.response }]);
}
catch (error) {
console.error("Error processing prompt:", error);
results.push([{ score: 0, output: undefined }]);
}
}
return results;
}
async ollamaQuery(prompt, options = {}) {
const requestBody = {
model: this.config.model,
prompt,
temperature: options.temperature ?? this.config.temperature,
stream: false,
format: this.config.structuredOutputFormat,
num_predict: options.maxDecodeSteps ?? this.config.maxTokens ?? 2048,
};
try {
const response = await axios_1.default.post(`${this.config.modelUrl}/api/generate`, requestBody, {
headers: { "Content-Type": "application/json" },
timeout: 30000,
});
return response.data;
}
catch (error) {
if (axios_1.default.isAxiosError(error)) {
throw new Error(`Ollama API error: ${error.response?.data?.error || error.message}`);
}
throw error;
}
}
}
exports.OllamaLanguageModel = OllamaLanguageModel;
//# sourceMappingURL=inference.js.map