📄 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
}