// hf-space/src/lib/knowledge/embedding.ts // 使用 Google Gemini Embedding API(gemini-embedding-001,1536 维输出) // 配置通过 ConfigManager 动态获取 import { getConfig } from '../config-manager'; let currentKeyIndex = 0; /** 获取 Google API 密钥列表(动态) */ async function getGoogleApiKeys(): Promise { const key1 = await getConfig('GOOGLE_EMBEDDING_KEY_1') || await getConfig('GOOGLE_AI_API_KEY'); const key2 = await getConfig('GOOGLE_EMBEDDING_KEY_2') || await getConfig('GOOGLE_AI_API_KEY_2'); return [key1, key2].filter(Boolean); } /** 轮换获取下一个密钥 */ async function getNextKey(): Promise { const keys = await getGoogleApiKeys(); if (keys.length === 0) throw new Error('No Google API keys configured'); const key = keys[currentKeyIndex % keys.length]; currentKeyIndex++; return key; } /** Google Gemini embedContent API 调用(单条,带超时) */ async function callGoogleEmbedding(text: string): Promise { const key = await getNextKey(); const model = await getConfig('GOOGLE_EMBEDDING_MODEL') || 'gemini-embedding-001'; const apiBase = await getConfig('GOOGLE_API_BASE') || 'https://generativelanguage.googleapis.com'; const url = `${apiBase}/v1beta/models/${model}:embedContent?key=${key}`; const controller = new AbortController(); const timeoutId = setTimeout(() => controller.abort(), 30000); try { const response = await fetch(url, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ model: `models/${model}`, content: { parts: [{ text }] }, outputDimensionality: 1536, }), signal: controller.signal, }); if (!response.ok) { const err = await response.text(); throw new Error(`Google Embedding API error: ${response.status} ${err}`); } const data = (await response.json()) as { embedding: { values: number[] } }; return data.embedding.values; } catch (err: unknown) { if (err instanceof Error && err.name === 'AbortError') { throw new Error('Google Embedding API timeout (30s)'); } throw err; } finally { clearTimeout(timeoutId); } } /** 批量向量化(逐条调用,轮换 key 分散限额) */ export async function embedBatch(texts: string[], retries = 3): Promise { const results: number[][] = []; for (const text of texts) { for (let attempt = 0; attempt < retries; attempt++) { try { const embedding = await callGoogleEmbedding(text); results.push(embedding); break; } catch (err) { if (attempt === retries - 1) throw err; // 429 限流退避 await new Promise((r) => setTimeout(r, 2000 * (attempt + 1))); } } } return results; } /** 计算 token 消耗 (用于成本追踪) */ export function estimateEmbeddingTokens(texts: string[]): number { return texts.reduce((total, text) => { const chineseChars = (text.match(/[一-鿿]/g) || []).length; const otherChars = text.length - chineseChars; return total + Math.ceil(chineseChars / 1.5 + otherChars / 4); }, 0); }