📄 embedding.ts • 5125 bytes
/**
* CmdCode 向量记忆系统 - Embedding 服务
* 使用加密存储的密钥
*/
import { sha256 } from './utils'
import {
loadMemorySearchConfig, DEFAULT_MEMORY_SEARCH,
getEmbeddingApiKey, rotateEmbeddingApiKey, getEmbeddingKeyPoolStatus,
resetEmbeddingKeyPool
} from '../apikeys.js'
import { t } from '../i18n.js'
/** 运行时读取配置(从加密存储或密钥池) */
export function loadConfig() {
// 从密钥池获取(不再支持环境变量)
const poolKey = getEmbeddingApiKey()
if (poolKey) {
return {
key: poolKey.apiKey, // 直接是 string 类型
baseUrl: DEFAULT_MEMORY_SEARCH.baseUrl,
model: DEFAULT_MEMORY_SEARCH.model,
source: poolKey.name
}
}
// 无密钥
return {
key: '',
baseUrl: DEFAULT_MEMORY_SEARCH.baseUrl,
model: DEFAULT_MEMORY_SEARCH.model,
source: 'none'
}
}
// 缓存
const embeddingCache = new Map<string, number[]>()
const CACHE_MAX_SIZE = 500
// 请求队列(控制并发)
let requestQueue: (() => void)[] = []
let activeRequests = 0
const MAX_CONCURRENT = 3
/** 获取 Embedding(带缓存和队列) */
export async function getEmbedding(text: string): Promise<number[]> {
// 运行时读取配置
const config = loadConfig()
if (!config.key) {
// 无嵌入密钥时返回空数组(内存搜索将被跳过)
console.warn('embedding: ARK_API_KEY 未配置,跳过向量搜索')
return []
}
// 检查缓存
const hash = sha256(text)
if (embeddingCache.has(hash)) {
return embeddingCache.get(hash)!
}
// 请求队列
return new Promise((resolve, reject) => {
const execute = async () => {
try {
activeRequests++
const embedding = await fetchEmbedding(text, config.key, config.baseUrl)
activeRequests--
// 存入缓存
if (embeddingCache.size >= CACHE_MAX_SIZE) {
const firstKey = embeddingCache.keys().next().value
embeddingCache.delete(firstKey)
}
embeddingCache.set(hash, embedding)
resolve(embedding)
processQueue()
} catch (e) {
activeRequests--
reject(e)
processQueue()
}
}
const processQueue = () => {
if (requestQueue.length > 0 && activeRequests < MAX_CONCURRENT) {
const next = requestQueue.shift()
if (next) next()
}
}
if (activeRequests < MAX_CONCURRENT) {
execute()
} else {
requestQueue.push(execute)
}
})
}
/** 批量获取 Embedding */
export async function getEmbeddings(texts: string[]): Promise<number[][]> {
const results: number[][] = []
for (const text of texts) {
try {
const emb = await getEmbedding(text)
results.push(emb)
} catch (e) {
console.error(t('error.embedding'), e)
results.push([])
}
}
return results
}
/** 直接调用 API(支持 429 自动切换密钥) */
async function fetchEmbedding(text: string, apiKey?: string, baseUrl?: string): Promise<number[]> {
const config = loadConfig()
if (!config.key) {
throw new Error('火山引擎 Embedding API 密钥池已耗尽或密钥库未解锁,请使用 /keypool add 命令添加新密钥')
}
const apiKeyFinal = apiKey || config.key
const apiBaseUrl = baseUrl || config.baseUrl || DEFAULT_MEMORY_SEARCH.baseUrl
const model = config.model || DEFAULT_MEMORY_SEARCH.model
const response = await fetch(apiBaseUrl + '/embeddings', {
method: 'POST',
headers: {
'Content-Type': 'application/json',
'Authorization': `Bearer ${apiKeyFinal}`
},
body: JSON.stringify({
model: model,
input: text.substring(0, 2000)
})
})
if (!response.ok) {
const err = await response.text()
// 检测 429 配额用尽,自动切换密钥
if (response.status === 429 || err.includes('quota') || err.includes('rate limit') || err.includes('exceeded')) {
const poolStatus = getEmbeddingKeyPoolStatus()
if (poolStatus.remaining > 0) {
const nextKey = rotateEmbeddingApiKey()
if (nextKey) {
console.log(` \x1b[33m⚠️ ${t('error.embedding_ratelimit', {name: nextKey.name})}\x1b[0m`)
return fetchEmbedding(text, nextKey.apiKey, baseUrl)
}
}
// 所有密钥都429→重置池,从头再轮
console.log(` \x1b[33m⚠️ 全部 ${poolStatus.total} 个 Embedding 密钥均触发限流,重置池后重新尝试\x1b[0m`)
resetEmbeddingKeyPool()
const nextKey = rotateEmbeddingApiKey()
if (nextKey) {
return fetchEmbedding(text, nextKey.apiKey, baseUrl)
}
throw new Error(`火山引擎 Embedding API 密钥池已耗尽(共 ${poolStatus.total} 个),请使用 /keypool add 命令添加新密钥`)
}
throw new Error(`Embedding API 错误: ${response.status} - ${err}`)
}
const data = await response.json() as any
return data.data[0].embedding
}
/** 清除缓存 */
export function clearCache(): void {
embeddingCache.clear()
}
/** 获取缓存大小 */
export function getCacheSize(): number {
return embeddingCache.size
}