@thecodingwhale/cv-processor
Version:
CV Processor to extract structured data from PDF resumes using TypeScript
226 lines (224 loc) • 8.98 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.GrokAIProvider = void 0;
const jsonrepair_1 = require("jsonrepair");
const openai_1 = require("openai");
const data_1 = require("../utils/data");
const GROK_AI_PRICING = {
'grok-3': { input: 0.003, output: 0.015 },
'grok-2-vision-1212': { input: 0.002, output: 0.01 },
'grok-2': { input: 0.002, output: 0.01 },
'grok-1': { input: 0.0001, output: 0.0002 }, // Older model with lower pricing
// Default
default: { input: 0.002, output: 0.01 },
};
/**
* Models that support structured output (response_format)
*/
const MODELS_WITH_STRUCTURED_OUTPUT = [
'grok-2-vision-1212',
'grok-2',
'grok-beta',
'grok-vision-beta',
];
class GrokAIProvider {
constructor(config) {
this.config = config;
console.log(`[GrokAIProvider] Initializing with model: ${config.model || 'grok-2-vision-1212'}`);
this.client = new openai_1.OpenAI({
apiKey: config.apiKey,
baseURL: 'https://api.x.ai/v1',
});
}
/**
* Check if the current model supports structured output
*/
supportsStructuredOutput(model) {
return MODELS_WITH_STRUCTURED_OUTPUT.some((supportedModel) => model.toLowerCase().includes(supportedModel.toLowerCase()));
}
/**
* Calculate estimated cost based on token usage and model
*/
calculateCost(promptTokens, completionTokens, model) {
// First try to match by specific model name
let pricing = GROK_AI_PRICING[model];
// If not found, try to match by partial model name
if (!pricing) {
const matchingKey = Object.keys(GROK_AI_PRICING).find((key) => model.toLowerCase().includes(key.toLowerCase()));
pricing = matchingKey
? GROK_AI_PRICING[matchingKey]
: GROK_AI_PRICING['default'];
}
const inputCost = (promptTokens / 1000) * pricing.input;
const outputCost = (completionTokens / 1000) * pricing.output;
return inputCost + outputCost;
}
/**
* Estimate token count based on text content
*/
estimateTokenCount(text) {
// Simple estimation: ~4 characters per token for English text
return Math.ceil(text.length / 4);
}
async extractStructuredDataFromImages(imageUrls, dataSchema, instructions) {
try {
const prompt = `
${instructions}
Extract information from the following document according to this JSON schema:
${JSON.stringify(dataSchema, null, 2)}
Your response should be valid JSON that matches this schema.
`;
// Check if the model supports vision capabilities
const modelName = this.config.model || 'grok-2-vision-1212';
// Create messages with the images
const messages = [
{
role: 'system',
content: prompt,
},
{
role: 'user',
content: [
{
type: 'text',
text: 'Please analyze this document:',
},
...imageUrls.map((imageUrl) => ({
type: 'image_url',
image_url: {
url: imageUrl,
},
})),
],
},
];
// Prepare the completion request
const completionRequest = {
model: modelName,
messages: messages,
};
// Only add response_format if the model supports it
if (this.supportsStructuredOutput(modelName)) {
completionRequest.response_format = { type: 'json_object' };
}
const completion = await this.client.chat.completions.create(completionRequest);
const responseText = completion.choices[0]?.message?.content || '{}';
// Extract token usage information
const promptTokens = completion.usage?.prompt_tokens ||
this.estimateTokenCount(prompt + JSON.stringify(imageUrls));
const completionTokens = completion.usage?.completion_tokens ||
this.estimateTokenCount(responseText);
const totalTokens = completion.usage?.total_tokens || promptTokens + completionTokens;
// Calculate estimated cost
const estimatedCost = this.calculateCost(promptTokens, completionTokens, modelName);
// Create token usage object
const tokenUsage = {
promptTokens,
completionTokens,
totalTokens,
estimatedCost,
};
try {
let fixedJson;
try {
fixedJson = (0, jsonrepair_1.jsonrepair)(responseText);
}
catch (err) {
try {
fixedJson = (0, jsonrepair_1.jsonrepair)(responseText);
}
catch (err) {
console.error('❌ Could not repair JSON:', err);
throw new Error(`AI returned invalid JSON: ${err}`);
}
}
const parsedJson = JSON.parse(fixedJson);
return {
...(0, data_1.replaceUUIDv4Placeholders)(parsedJson),
tokenUsage,
};
}
catch (jsonError) {
console.error('Error parsing JSON from OpenAI response:', jsonError);
throw jsonError;
}
}
catch (error) {
console.error('Error extracting structured data with Grok AI:', error);
throw error;
}
}
async extractStructuredDataFromText(texts, dataSchema, instructions, categories) {
try {
const modelName = this.config.model || 'grok-2-vision-1212';
const prompt = `
${instructions}
Extract information from the following text according to this JSON schema:
${JSON.stringify(dataSchema, null, 2)}
Your response should be valid JSON that matches this schema.
Text content:
${texts.join('\n\n')}
`;
// Prepare the completion request
const completionRequest = {
model: modelName,
messages: [
{
role: 'system',
content: prompt,
},
],
};
// Only add response_format if the model supports it
if (this.supportsStructuredOutput(modelName)) {
completionRequest.response_format = { type: 'json_object' };
}
const completion = await this.client.chat.completions.create(completionRequest);
const responseText = completion.choices[0]?.message?.content || '{}';
// Extract token usage information
const promptTokens = completion.usage?.prompt_tokens || this.estimateTokenCount(prompt);
const completionTokens = completion.usage?.completion_tokens ||
this.estimateTokenCount(responseText);
const totalTokens = completion.usage?.total_tokens || promptTokens + completionTokens;
// Calculate estimated cost
const estimatedCost = this.calculateCost(promptTokens, completionTokens, modelName);
// Create token usage object
const tokenUsage = {
promptTokens,
completionTokens,
totalTokens,
estimatedCost,
};
try {
let fixedJson;
try {
fixedJson = (0, jsonrepair_1.jsonrepair)(responseText);
}
catch (err) {
console.error('❌ Could not repair JSON:', err);
throw new Error(`AI returned invalid JSON: ${err}`);
}
const parsedJson = JSON.parse(fixedJson);
return {
...(0, data_1.replaceUUIDv4Placeholders)(parsedJson),
tokenUsage,
};
}
catch (jsonError) {
console.error('Error parsing JSON from Grok AI response:', jsonError);
throw jsonError;
}
}
catch (error) {
console.error('Error extracting structured data with Grok AI:', error);
throw error;
}
}
getModelInfo() {
return {
provider: 'grok',
model: this.config.model || 'grok-2-vision-1212',
};
}
}
exports.GrokAIProvider = GrokAIProvider;