68 lines
2.6 KiB
JavaScript
68 lines
2.6 KiB
JavaScript
'use strict';
|
|
|
|
const { Pool } = require('pg');
|
|
|
|
const cache = new Map();
|
|
const CACHE_MAX = 200;
|
|
|
|
function createEmbeddingClient(cfg, opts = {}) {
|
|
const pool = new Pool({ connectionString: opts.pgUrl || cfg.ragPgUrl, max: 2, connectionTimeoutMillis: 5000 });
|
|
const model = cfg.embModel || 'text-embedding-v4';
|
|
const dim = cfg.embDim || 1024;
|
|
const baseUrl = String(cfg.embBaseUrl || '').replace(/\/+$/, '');
|
|
const apiKey = cfg.embApiKey;
|
|
const translateModel = cfg.embTranslateModel || 'qwen-turbo';
|
|
|
|
async function embedQuery(text) {
|
|
const key = String(text || '').slice(0, 200);
|
|
if (cache.has(key)) return cache.get(key);
|
|
if (!baseUrl || !apiKey) throw new Error('Embedding provider is not configured');
|
|
const res = await fetch(`${baseUrl}/embeddings`, {
|
|
method: 'POST',
|
|
headers: { Authorization: `Bearer ${apiKey}`, 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({ model, input: text, encoding_format: 'float', dimensions: dim }),
|
|
signal: AbortSignal.timeout(15000),
|
|
});
|
|
const json = await res.json();
|
|
if (!res.ok || !json.data || !json.data[0] || !json.data[0].embedding) {
|
|
throw new Error('embedQuery failed: ' + JSON.stringify(json).slice(0, 200));
|
|
}
|
|
const vec = json.data[0].embedding;
|
|
if (cache.size >= CACHE_MAX) cache.delete(cache.keys().next().value);
|
|
cache.set(key, vec);
|
|
return vec;
|
|
}
|
|
|
|
function vecStr(v) {
|
|
return '[' + v.join(',') + ']';
|
|
}
|
|
|
|
async function translateQueryToEnglish(query) {
|
|
try {
|
|
const res = await fetch(`${baseUrl}/chat/completions`, {
|
|
method: 'POST',
|
|
headers: { Authorization: `Bearer ${apiKey}`, 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({
|
|
model: translateModel,
|
|
max_tokens: 80,
|
|
messages: [
|
|
{ role: 'system', content: 'Translate the user message into a concise English search query for a B2B jewelry wholesaler FAQ. Output ONLY the English translation, no explanation, no quotes.' },
|
|
{ role: 'user', content: query },
|
|
],
|
|
}),
|
|
signal: AbortSignal.timeout(15000),
|
|
});
|
|
const json = await res.json();
|
|
const out = (((json.choices || [])[0] || {}).message || {}).content || '';
|
|
return String(out).trim().replace(/^["']|["']$/g, '');
|
|
} catch (err) {
|
|
console.error(JSON.stringify({ ts: new Date().toISOString(), warn: `kb translate failed: ${err.message}` }));
|
|
return null;
|
|
}
|
|
}
|
|
|
|
return { pool, embedQuery, vecStr, translateQueryToEnglish, model, dim };
|
|
}
|
|
|
|
module.exports = { createEmbeddingClient };
|