KnowledgeBase/server/src/knowledge-base/embedding.service.ts

import { Injectable, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { ModelConfigService } from '../model-config/model-config.service';

export interface EmbeddingResponse {
  data: Array<{
    embedding: number[];
    index: number;
  }>;
  model: string;
  usage: {
    prompt_tokens: number;
    total_tokens: number;
  };
}

@Injectable()
export class EmbeddingService {
  private readonly logger = new Logger(EmbeddingService.name);
  private readonly defaultDimensions: number;

  constructor(
    private modelConfigService: ModelConfigService,
    private configService: ConfigService,
  ) {
    this.defaultDimensions = parseInt(
      this.configService.get<string>('DEFAULT_VECTOR_DIMENSIONS', '2560'),
    );
    this.logger.log(`デフォルトのベクトル次元が ${this.defaultDimensions} に設定されました`);
  }

  async getEmbeddings(
    texts: string[],
    userId: string,
    embeddingModelConfigId: string,
    tenantId?: string,
  ): Promise<number[][]> {
    this.logger.log(`${texts.length} 個のテキストに対して埋め込みベクトルを生成しています`);

    const modelConfig = await this.modelConfigService.findOne(
      embeddingModelConfigId,
      userId,
      tenantId || 'default',
    );
    if (!modelConfig || modelConfig.type !== 'embedding') {
      throw new Error(`埋め込みモデル設定 ${embeddingModelConfigId} が見つかりません`);
    }

    if (modelConfig.isEnabled === false) {
      throw new Error(`モデル ${modelConfig.name} は無効化されているため、埋め込みベクトルを生成できません`);
    }

    // APIキーはオプションです - ローカルモデルを許可します

    if (!modelConfig.baseUrl) {
      throw new Error(`モデル ${modelConfig.name} に baseUrl が設定されていません`);
    }

    // モデル名に基づいて最大バッチサイズを決定
    const maxBatchSize = this.getMaxBatchSizeForModel(modelConfig.modelId, modelConfig.maxBatchSize);

    // バッチサイズが制限を超える場合は分割して処理
    if (texts.length > maxBatchSize) {
      this.logger.log(
        `テキスト数 ${texts.length} がモデルのバッチ制限 ${maxBatchSize} を超えているため、分割処理します`
      );

      const allEmbeddings: number[][] = [];

      for (let i = 0; i < texts.length; i += maxBatchSize) {
        const batch = texts.slice(i, i + maxBatchSize);
        const batchEmbeddings = await this.getEmbeddingsForBatch(
          batch,
          userId,
          modelConfig,
          maxBatchSize
        );

        allEmbeddings.push(...batchEmbeddings);

        // APIレート制限対策のため、短い間隔で待機
        if (i + maxBatchSize < texts.length) {
          await new Promise(resolve => setTimeout(resolve, 100)); // 100ms待機
        }
      }

      return allEmbeddings;
    } else {
      // 通常処理（バッチサイズ以内）
      return await this.getEmbeddingsForBatch(
        texts,
        userId,
        modelConfig,
        maxBatchSize
      );
    }
  }

  /**
   * モデルIDに基づいて最大バッチサイズを決定
   */
  private getMaxBatchSizeForModel(modelId: string, configuredMaxBatchSize?: number): number {
    // モデル固有のバッチサイズ制限
    if (modelId.includes('text-embedding-004') || modelId.includes('text-embedding-v4') ||
      modelId.includes('text-embedding-ada-002')) {
      return Math.min(10, configuredMaxBatchSize || 100); // Googleの場合は10を上限
    } else if (modelId.includes('text-embedding-3') || modelId.includes('text-embedding-003')) {
      return Math.min(2048, configuredMaxBatchSize || 2048); // OpenAI v3は2048が上限
    } else {
      // デフォルトでは設定された最大バッチサイズか100の小さい方
      return Math.min(configuredMaxBatchSize || 100, 100);
    }
  }

  /**
   * 単一バッチの埋め込み処理
   */
  private async getEmbeddingsForBatch(
    texts: string[],
    userId: string,
    modelConfig: any,
    maxBatchSize: number,
  ): Promise<number[][]> {
    const apiUrl = modelConfig.baseUrl.endsWith('/embeddings')
      ? modelConfig.baseUrl
      : `${modelConfig.baseUrl}/embeddings`;

    let lastError;
    const MAX_RETRIES = 3;

    for (let attempt = 1; attempt <= MAX_RETRIES; attempt++) {
      try {
        const controller = new AbortController();
        const timeoutId = setTimeout(() => {
          controller.abort();
          this.logger.error(`Embedding API timeout after 60s: ${apiUrl}`);
        }, 60000); // 60s timeout

        this.logger.log(`[モデル呼び出し] タイプ: Embedding, モデル: ${modelConfig.name} (${modelConfig.modelId}), ユーザー: ${userId}, テキスト数: ${texts.length}`);
        this.logger.log(`埋め込み API を呼び出し中 (試行 ${attempt}/${MAX_RETRIES}): ${apiUrl}`);

        let response;
        try {
          response = await fetch(apiUrl, {
            method: 'POST',
            headers: {
              'Content-Type': 'application/json',
              Authorization: `Bearer ${modelConfig.apiKey}`,
            },
            body: JSON.stringify({
              encoding_format: 'float',
              input: texts,
              model: modelConfig.modelId,
            }),
            signal: controller.signal,
          });
        } finally {
          clearTimeout(timeoutId);
        }

        if (!response.ok) {
          const errorText = await response.text();

          // バッチサイズ制限エラーを検出
          if (errorText.includes('batch size is invalid') || errorText.includes('batch_size') ||
            errorText.includes('invalid') || errorText.includes('larger than')) {
            this.logger.warn(
              `バッチサイズ制限エラーが検出されました。バッチサイズを半分に分割して再試行します: ${maxBatchSize} -> ${Math.floor(maxBatchSize / 2)}`
            );

            // バッチをさらに小さな単位に分割して再試行
            if (texts.length > 1) {
              const midPoint = Math.floor(texts.length / 2);
              const firstHalf = texts.slice(0, midPoint);
              const secondHalf = texts.slice(midPoint);

              const firstResult = await this.getEmbeddingsForBatch(firstHalf, userId, modelConfig, Math.floor(maxBatchSize / 2));
              const secondResult = await this.getEmbeddingsForBatch(secondHalf, userId, modelConfig, Math.floor(maxBatchSize / 2));

              return [...firstResult, ...secondResult];
            }
          }

          // コンテキスト長の過剰エラーを検出
          if (errorText.includes('context length') || errorText.includes('exceeds')) {
            const avgLength = texts.reduce((s, t) => s + t.length, 0) / texts.length;
            const totalLength = texts.reduce((s, t) => s + t.length, 0);
            this.logger.error(
              `テキスト長が制限を超過しました: 入力 ${texts.length} 個のテキスト、` +
              `総計 ${totalLength} 文字、平均 ${Math.round(avgLength)} 文字、` +
              `モデル制限: ${modelConfig.maxInputTokens || 8192} tokens`
            );
            throw new Error(
              `テキスト長がモデルの制限を超えています。` +
              `現在: ${texts.length} 個のテキストで計 ${totalLength} 文字、` +
              `モデル制限: ${modelConfig.maxInputTokens || 8192} tokens。` +
              `アドバイス: チャンクサイズまたはバッチサイズを小さくしてください`
            );
          }

          // 429 (Too Many Requests) または 5xx (Server Error) の場合は再試行
          if (response.status === 429 || response.status >= 500) {
            this.logger.warn(`埋め込み API で一時的なエラーが発生しました (${response.status}): ${errorText}`);
            throw new Error(`API Error ${response.status}: ${errorText}`);
          }

          this.logger.error(`埋め込み API エラーの詳細: ${errorText}`);
          this.logger.error(`リクエストパラメータ: model=${modelConfig.modelId}, inputLength=${texts[0]?.length}`);
          throw new Error(`埋め込み API の呼び出しに失敗しました: ${response.statusText} - ${errorText}`);
        }

        const data: EmbeddingResponse = await response.json();
        const embeddings = data.data.map((item) => item.embedding);

        // 実際のレスポンスから次元を取得
        const actualDimensions = embeddings[0]?.length || this.defaultDimensions;
        this.logger.log(
          `${modelConfig.name} から ${embeddings.length} 個の埋め込みベクトルを取得しました。次元: ${actualDimensions}`,
        );

        return embeddings;
      } catch (error) {
        lastError = error;

        // 最後のアテンプトでなく、エラーが一時的と思われる場合（または堅牢性のために全て）は、待機後に再試行
        if (attempt < MAX_RETRIES) {
          const delay = Math.pow(2, attempt - 1) * 1000; // 1s, 2s, 4s
          this.logger.warn(`埋め込みリクエストが失敗しました。${delay}ms 後に再試行します: ${error.message}`);
          await new Promise(resolve => setTimeout(resolve, delay));
          continue;
        }
      }
    }

    throw lastError;
  }

  private getEstimatedDimensions(modelId: string): number {
    // 使用环境变量的默认维度
    return this.defaultDimensions;
  }
}