src/lib/knowledge/embedding.ts

// 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<string[]> {
  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<string> {
  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<number[]> {
  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<number[][]> {
  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);
}