118 lines
3.4 KiB
JavaScript
118 lines
3.4 KiB
JavaScript
/**
|
|
* OpenAI Provider - Handles embedding generation
|
|
*/
|
|
|
|
const EMBEDDING_MODEL = 'text-embedding-3-small';
|
|
const EMBEDDING_DIMENSIONS = 1536;
|
|
const MAX_BATCH_SIZE = 2048;
|
|
|
|
class OpenAIProvider {
|
|
constructor({ apiKey, baseUrl = 'https://api.openai.com/v1', timeoutMs = 60000 }) {
|
|
if (!apiKey) {
|
|
throw new Error('OpenAI API key is required');
|
|
}
|
|
this.apiKey = apiKey;
|
|
this.baseUrl = baseUrl;
|
|
this.timeoutMs = timeoutMs;
|
|
}
|
|
|
|
/**
|
|
* Generate embeddings for one or more texts
|
|
* @param {string|string[]} input - Text or array of texts
|
|
* @param {Object} options
|
|
* @returns {Promise<{embeddings: number[][], usage: Object, model: string, latencyMs: number}>}
|
|
*/
|
|
async embed(input, options = {}) {
|
|
const texts = Array.isArray(input) ? input : [input];
|
|
const model = options.model || EMBEDDING_MODEL;
|
|
const dimensions = options.dimensions || EMBEDDING_DIMENSIONS;
|
|
const timeoutMs = options.timeoutMs || this.timeoutMs;
|
|
|
|
if (texts.length > MAX_BATCH_SIZE) {
|
|
throw new Error(`Batch size ${texts.length} exceeds max of ${MAX_BATCH_SIZE}`);
|
|
}
|
|
|
|
const started = Date.now();
|
|
|
|
// Clean and truncate input texts
|
|
const cleanedTexts = texts.map(t =>
|
|
(t || '').replace(/\n+/g, ' ').trim().substring(0, 8000)
|
|
);
|
|
|
|
const body = {
|
|
input: cleanedTexts,
|
|
model,
|
|
encoding_format: 'float'
|
|
};
|
|
|
|
// Only embedding-3 models support dimensions parameter
|
|
if (model.includes('embedding-3')) {
|
|
body.dimensions = dimensions;
|
|
}
|
|
|
|
const response = await this._makeRequest('embeddings', body, timeoutMs);
|
|
|
|
// Sort by index to ensure order matches input
|
|
const sortedData = response.data.sort((a, b) => a.index - b.index);
|
|
|
|
return {
|
|
embeddings: sortedData.map(item => item.embedding),
|
|
usage: {
|
|
promptTokens: response.usage?.prompt_tokens || 0,
|
|
totalTokens: response.usage?.total_tokens || 0
|
|
},
|
|
model: response.model || model,
|
|
latencyMs: Date.now() - started
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Generator for processing large batches in chunks
|
|
*/
|
|
async *embedBatchChunked(texts, options = {}) {
|
|
const batchSize = Math.min(options.batchSize || 100, MAX_BATCH_SIZE);
|
|
|
|
for (let i = 0; i < texts.length; i += batchSize) {
|
|
const chunk = texts.slice(i, i + batchSize);
|
|
const result = await this.embed(chunk, options);
|
|
|
|
yield {
|
|
embeddings: result.embeddings,
|
|
startIndex: i,
|
|
endIndex: i + chunk.length,
|
|
usage: result.usage,
|
|
model: result.model,
|
|
latencyMs: result.latencyMs
|
|
};
|
|
}
|
|
}
|
|
|
|
async _makeRequest(endpoint, body, timeoutMs) {
|
|
const controller = new AbortController();
|
|
const timeout = setTimeout(() => controller.abort(), timeoutMs);
|
|
|
|
try {
|
|
const response = await fetch(`${this.baseUrl}/${endpoint}`, {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
'Authorization': `Bearer ${this.apiKey}`
|
|
},
|
|
body: JSON.stringify(body),
|
|
signal: controller.signal
|
|
});
|
|
|
|
if (!response.ok) {
|
|
const error = await response.json().catch(() => ({}));
|
|
throw new Error(error.error?.message || `OpenAI API error: ${response.status}`);
|
|
}
|
|
|
|
return response.json();
|
|
} finally {
|
|
clearTimeout(timeout);
|
|
}
|
|
}
|
|
}
|
|
|
|
module.exports = { OpenAIProvider, EMBEDDING_MODEL, EMBEDDING_DIMENSIONS };
|