@thecodingwhale/cv-processor
Version:
CV Processor to extract structured data from PDF resumes using TypeScript
263 lines (259 loc) • 10.9 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.GeminiAIProvider = void 0;
const genai_1 = require("@google/genai");
const jsonrepair_1 = require("jsonrepair");
const data_1 = require("../utils/data");
const GEMINI_PRICING = {
// Gemini 2.5 models
'gemini-2.5-flash': { input: 0.00015, output: 0.0006 },
'gemini-2.5-pro': { input: 0.00125, output: 0.01 },
// Gemini 2.0 models
'gemini-2.0-flash': { input: 0.0001, output: 0.0004 },
'gemini-2.0-flash-lite': { input: 0.000075, output: 0.0003 },
// Gemini 1.5 models
'gemini-1.5-pro': { input: 0.00125, output: 0.005 },
'gemini-1.5-flash': { input: 0.000075, output: 0.0003 },
'gemini-1.5-flash-8b': { input: 0.0000375, output: 0.00015 },
// Legacy models
'gemini-pro': { input: 0.00125, output: 0.00375 },
// Default fallback pricing (using Gemini 1.5 Pro as baseline)
default: { input: 0.00125, output: 0.005 },
};
class GeminiAIProvider {
constructor(config) {
this.config = config;
this.ai = new genai_1.GoogleGenAI({ apiKey: config.apiKey });
}
/**
* Calculate estimated cost based on token usage and model
*/
calculateCost(promptTokens, completionTokens, model) {
const pricing = GEMINI_PRICING[model] || GEMINI_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);
}
/**
* Convert image URLs to proper content parts for the new API
*/
async createImageParts(imageUrls) {
const imageParts = [];
for (const imageUrl of imageUrls) {
try {
if (imageUrl.startsWith('data:image/')) {
// Handle data URLs (base64-encoded images)
const [mimeTypePart, base64Data] = imageUrl.split(',');
const mimeType = mimeTypePart.split(':')[1].split(';')[0];
imageParts.push({
inlineData: {
mimeType: mimeType,
data: base64Data,
},
});
}
else if (imageUrl.startsWith('http://') ||
imageUrl.startsWith('https://')) {
// Handle web URLs by fetching and converting to base64
const response = await fetch(imageUrl);
if (!response.ok) {
console.warn(`Failed to fetch image from ${imageUrl}: ${response.statusText}`);
continue;
}
const arrayBuffer = await response.arrayBuffer();
const base64Data = Buffer.from(arrayBuffer).toString('base64');
// Determine MIME type from response headers or URL extension
let mimeType = response.headers.get('content-type') || 'image/jpeg';
if (!mimeType.startsWith('image/')) {
// Fallback based on URL extension
if (imageUrl.toLowerCase().includes('.png')) {
mimeType = 'image/png';
}
else if (imageUrl.toLowerCase().includes('.webp')) {
mimeType = 'image/webp';
}
else {
mimeType = 'image/jpeg';
}
}
imageParts.push({
inlineData: {
mimeType: mimeType,
data: base64Data,
},
});
}
else if (imageUrl.startsWith('gs://')) {
// Handle Google Cloud Storage URLs
imageParts.push({
fileData: {
mimeType: 'image/jpeg', // Default, could be improved with better detection
fileUri: imageUrl,
},
});
}
else {
console.warn(`Unsupported image URL format: ${imageUrl}`);
}
}
catch (error) {
console.error(`Error processing image URL ${imageUrl}:`, error);
}
}
return imageParts;
}
async extractStructuredDataFromImages(imageUrls, dataSchema, instructions) {
try {
const model = this.config.model || 'gemini-1.5-pro';
// Create proper image content parts
const imageParts = await this.createImageParts(imageUrls);
if (imageParts.length === 0) {
throw new Error('No valid images could be processed from the provided URLs');
}
// Create the content array with text and images
const contents = [
{ text: instructions },
{
text: `Extract information according to this JSON schema: ${JSON.stringify(dataSchema, null, 2)}`,
},
{
text: 'Your response should be valid JSON that matches this schema.',
},
...imageParts,
];
// Set topK to 40 if using gemini-1.5-flash-8b model which has a limitation
const topK = model === 'gemini-1.5-flash-8b' ? 40 : 50;
const result = await this.ai.models.generateContent({
model: model,
contents: contents,
config: {
temperature: this.config.temperature || 0,
maxOutputTokens: this.config.maxTokens || 8192,
topP: 1,
topK: topK,
candidateCount: 1,
},
});
const responseText = result.text || '';
// Estimate token usage (the new API doesn't provide easy access to token counts)
const promptTokens = this.estimateTokenCount(instructions + JSON.stringify(dataSchema)) +
imageParts.length * 258; // 258 tokens per image
const completionTokens = this.estimateTokenCount(responseText);
const totalTokens = promptTokens + completionTokens;
// Calculate estimated cost
const estimatedCost = this.calculateCost(promptTokens, completionTokens, model);
// 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 Gemini response:', jsonError);
throw jsonError;
}
}
catch (error) {
console.error('Error extracting structured data with Gemini AI:', error);
throw error;
}
}
async extractStructuredDataFromText(texts, dataSchema, instructions, categories) {
try {
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')}
`;
const model = this.config.model || 'gemini-1.5-pro';
// Set topK to 40 if using gemini-1.5-flash-8b model which has a limitation
const topK = model === 'gemini-1.5-flash-8b' ? 40 : 50;
const result = await this.ai.models.generateContent({
model: model,
contents: prompt,
config: {
temperature: this.config.temperature || 0,
maxOutputTokens: this.config.maxTokens || 8192,
topP: 1,
topK: topK,
candidateCount: 1,
},
});
const responseText = result.text || '';
// Estimate token usage
const promptTokens = this.estimateTokenCount(prompt);
const completionTokens = this.estimateTokenCount(responseText);
const totalTokens = promptTokens + completionTokens;
// Calculate estimated cost
const estimatedCost = this.calculateCost(promptTokens, completionTokens, model);
// 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 Gemini response:', jsonError);
throw jsonError;
}
}
catch (error) {
console.error('Error extracting structured data with Gemini AI:', error);
throw error;
}
}
getModelInfo() {
return {
provider: 'gemini',
model: this.config.model || 'gemini-1.5-pro',
};
}
}
exports.GeminiAIProvider = GeminiAIProvider;