@just-every/ensemble
Version:
LLM provider abstraction layer with unified streaming interface
58 lines • 2.45 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.ensembleEmbed = ensembleEmbed;
const model_provider_js_1 = require("../model_providers/model_provider.cjs");
const model_data_js_1 = require("../data/model_data.cjs");
const EMBEDDING_TTL_MS = 1000 * 60 * 60;
const EMBEDDING_CACHE_MAX = 1000;
const embeddingCache = new Map();
async function ensembleEmbed(text, agent, options) {
const dimensions = options?.dimensions || 1536;
const cacheKey = `${agent.model || agent.modelClass}:${text}:${dimensions}`;
const cached = embeddingCache.get(cacheKey);
if (cached) {
if (Date.now() - cached.timestamp.getTime() < EMBEDDING_TTL_MS) {
return cached.embedding;
}
embeddingCache.delete(cacheKey);
}
const model = agent.model || 'text-embedding-3-small';
const provider = (0, model_provider_js_1.getModelProvider)(model);
if (!provider.createEmbedding) {
throw new Error(`Provider for model ${model} does not support embeddings`);
}
const modelInfo = (0, model_data_js_1.findModel)(model);
const inputTokenLimit = modelInfo?.features?.input_token_limit;
const MAX_CHARS_PER_CHUNK = inputTokenLimit ? inputTokenLimit * 4 * 0.9 : Infinity;
const needsChunking = inputTokenLimit && text.length > MAX_CHARS_PER_CHUNK;
let embedding;
if (needsChunking) {
const chunks = [];
for (let i = 0; i < text.length; i += MAX_CHARS_PER_CHUNK) {
chunks.push(text.slice(i, i + MAX_CHARS_PER_CHUNK));
}
const result = await provider.createEmbedding(chunks, model, { ...options, dimensions });
const embeddings = result;
embedding = new Array(dimensions).fill(0);
for (const vec of embeddings) {
for (let i = 0; i < dimensions; i++) {
embedding[i] += vec[i] / embeddings.length;
}
}
}
else {
const result = await provider.createEmbedding(text, model, { ...options, dimensions });
embedding = Array.isArray(result[0]) ? result[0] : result;
}
if (embeddingCache.size >= EMBEDDING_CACHE_MAX) {
const oldestKey = embeddingCache.keys().next().value;
if (oldestKey)
embeddingCache.delete(oldestKey);
}
embeddingCache.set(cacheKey, {
embedding,
timestamp: new Date(),
});
return embedding;
}
//# sourceMappingURL=ensemble_embed.js.map