KnowledgeBase/server/src/knowledge-base/chunk-config.service.ts

import { Injectable, Logger, BadRequestException } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { ModelConfigService } from '../model-config/model-config.service';
import { TenantService } from '../tenant/tenant.service';
// import { UserSettingService } from '../user-setting/user-setting.service';

/**
 * チャンク設定サービス
 * チャンクパラメータの検証と管理を担当し、モデルの制限や環境変数の設定に適合していることを確認します
 *
 * 制限の優先順位:
 * 1. 環境変数 (MAX_CHUNK_SIZE, MAX_OVERLAP_SIZE)
 * 2. データベース内のモデル設定 (maxInputTokens, maxBatchSize)
 * 3. デフォルト値
 */
import {
  DEFAULT_CHUNK_SIZE,
  MIN_CHUNK_SIZE,
  DEFAULT_CHUNK_OVERLAP,
  MIN_CHUNK_OVERLAP,
  DEFAULT_MAX_OVERLAP_RATIO,
  DEFAULT_MAX_BATCH_SIZE,
  DEFAULT_VECTOR_DIMENSIONS
} from '../common/constants';
import { I18nService } from '../i18n/i18n.service';

@Injectable()
export class ChunkConfigService {
  private readonly logger = new Logger(ChunkConfigService.name);

  // デフォルト設定
  private readonly DEFAULTS = {
    chunkSize: DEFAULT_CHUNK_SIZE,
    chunkOverlap: DEFAULT_CHUNK_OVERLAP,
    minChunkSize: MIN_CHUNK_SIZE,
    minChunkOverlap: MIN_CHUNK_OVERLAP,
    maxOverlapRatio: DEFAULT_MAX_OVERLAP_RATIO,  // 重なりはチャンクサイズの50%まで
    maxBatchSize: DEFAULT_MAX_BATCH_SIZE,    // デフォルトのバッチ制限
    expectedDimensions: DEFAULT_VECTOR_DIMENSIONS, // デフォルトのベクトル次元
  };

  // 環境変数で設定された上限（優先的に使用）
  private readonly envMaxChunkSize: number;
  private readonly envMaxOverlapSize: number;

  constructor(
    private configService: ConfigService,
    private modelConfigService: ModelConfigService,
    private i18nService: I18nService,
    private tenantService: TenantService,
  ) {
    // 環境変数からグローバルな上限設定を読み込む
    this.envMaxChunkSize = parseInt(
      this.configService.get<string>('MAX_CHUNK_SIZE', '8191')
    );
    this.envMaxOverlapSize = parseInt(
      this.configService.get<string>('MAX_OVERLAP_SIZE', '2000')
    );

    this.logger.log(
      `環境変数設定の上限: MAX_CHUNK_SIZE=${this.envMaxChunkSize}, MAX_OVERLAP_SIZE=${this.envMaxOverlapSize}`
    );
  }

  /**
   * モデルの制限設定を取得（データベースから読み込み）
   */
  async getModelLimits(modelId: string, userId: string, tenantId?: string): Promise<{
    maxInputTokens: number;
    maxBatchSize: number;
    expectedDimensions: number;
    providerName: string;
    isVectorModel: boolean;
  }> {
    const modelConfig = await this.modelConfigService.findOne(modelId, userId, tenantId || '');

    if (!modelConfig || modelConfig.type !== 'embedding') {
      throw new BadRequestException(this.i18nService.formatMessage('embeddingModelNotFound', { id: modelId }));
    }

    // データベースのフィールドから制限を取得し、デフォルト値で補完
    const maxInputTokens = modelConfig.maxInputTokens || this.envMaxChunkSize;
    const maxBatchSize = modelConfig.maxBatchSize || this.DEFAULTS.maxBatchSize;
    const expectedDimensions = modelConfig.dimensions || parseInt(this.configService.get('DEFAULT_VECTOR_DIMENSIONS', String(this.DEFAULTS.expectedDimensions)));
    const providerName = modelConfig.providerName || '不明';
    const isVectorModel = modelConfig.isVectorModel || false;

    this.logger.log(
      this.i18nService.formatMessage('configLoaded', { name: modelConfig.name, id: modelConfig.modelId }) + '\n' +
      `  - プロバイダー: ${providerName}\n` +
      `  - Token制限: ${maxInputTokens}\n` +
      `  - バッチ制限: ${maxBatchSize}\n` +
      `  - ベクトル次元: ${expectedDimensions}\n` +
      `  - ベクトルモデルか: ${isVectorModel}`,
    );

    return {
      maxInputTokens,
      maxBatchSize,
      expectedDimensions,
      providerName,
      isVectorModel,
    };
  }

  /**
   * チャンク設定を検証および修正
   * 優先順位: 環境変数の上限 > モデルの制限 > ユーザー設定
   */
  async validateChunkConfig(
    chunkSize: number,
    chunkOverlap: number,
    modelId: string,
    userId: string,
    tenantId?: string,
  ): Promise<{
    chunkSize: number;
    chunkOverlap: number;
    warnings: string[];
    effectiveMaxChunkSize: number;
    effectiveMaxOverlapSize: number;
  }> {
    const warnings: string[] = [];
    const limits = await this.getModelLimits(modelId, userId, tenantId);

    // 1. 最終的な上限を計算（環境変数とモデル制限の小さい方を選択）
    const effectiveMaxChunkSize = Math.min(
      this.envMaxChunkSize,
      limits.maxInputTokens,
    );

    const effectiveMaxOverlapSize = Math.min(
      this.envMaxOverlapSize,
      Math.floor(effectiveMaxChunkSize * this.DEFAULTS.maxOverlapRatio),
    );

    // 2. チャンクサイズの上限を検証
    if (chunkSize > effectiveMaxChunkSize) {
      const reason =
        this.envMaxChunkSize < limits.maxInputTokens
          ? `${this.i18nService.getMessage('environmentLimit')} ${this.envMaxChunkSize}`
          : `${this.i18nService.getMessage('modelLimit')} ${limits.maxInputTokens}`;

      warnings.push(
        this.i18nService.formatMessage('chunkOverflow', {
          size: chunkSize,
          max: effectiveMaxChunkSize,
          reason
        })
      );
      chunkSize = effectiveMaxChunkSize;
    }

    // 3. チャンクサイズの下限を検証
    if (chunkSize < this.DEFAULTS.minChunkSize) {
      warnings.push(
        this.i18nService.formatMessage('chunkUnderflow', {
          size: chunkSize,
          min: this.DEFAULTS.minChunkSize
        })
      );
      chunkSize = this.DEFAULTS.minChunkSize;
    }

    // 4. 重なりサイズの上限を検証（環境変数優先）
    if (chunkOverlap > effectiveMaxOverlapSize) {
      warnings.push(
        this.i18nService.formatMessage('overlapOverflow', {
          size: chunkOverlap,
          max: effectiveMaxOverlapSize
        })
      );
      chunkOverlap = effectiveMaxOverlapSize;
    }

    // 5. 重なりサイズがチャンクサイズの50%を超えないことを検証
    const maxOverlapByRatio = Math.floor(
      chunkSize * this.DEFAULTS.maxOverlapRatio,
    );
    if (chunkOverlap > maxOverlapByRatio) {
      warnings.push(
        this.i18nService.formatMessage('overlapRatioExceeded', {
          size: chunkOverlap,
          max: maxOverlapByRatio
        })
      );
      chunkOverlap = maxOverlapByRatio;
    }

    if (chunkOverlap < this.DEFAULTS.minChunkOverlap) {
      warnings.push(
        this.i18nService.formatMessage('overlapUnderflow', {
          size: chunkOverlap,
          min: this.DEFAULTS.minChunkOverlap
        })
      );
      chunkOverlap = this.DEFAULTS.minChunkOverlap;
    }

    // 6. バッチ処理の安全チェックを追加
    // バッチ処理時、複数のテキストの合計長がモデルの制限を超えないようにする必要があります
    const safetyMargin = 0.8; // 80% 安全マージン、バッチ処理のためにスペースを確保
    const safeChunkSize = Math.floor(effectiveMaxChunkSize * safetyMargin);

    if (chunkSize > safeChunkSize) {
      warnings.push(
        this.i18nService.formatMessage('batchOverflowWarning', {
          safeSize: safeChunkSize,
          size: chunkSize,
          percent: Math.round(safetyMargin * 100)
        })
      );
    }

    // 7. 推定チャンク数が妥当かチェック
    const estimatedChunkCount = this.estimateChunkCount(
      1000000, // 1MB のテキストを想定
      chunkSize,
    );

    if (estimatedChunkCount > 50000) {
      warnings.push(
        this.i18nService.formatMessage('estimatedChunkCountExcessive', { count: estimatedChunkCount })
      );
    }

    return {
      chunkSize,
      chunkOverlap,
      warnings,
      effectiveMaxChunkSize,
      effectiveMaxOverlapSize,
    };
  }

  /**
   * 推奨されるバッチサイズを取得
   */
  async getRecommendedBatchSize(
    modelId: string,
    userId: string,
    tenantId?: string,
    currentBatchSize: number = 100,
  ): Promise<number> {
    const limits = await this.getModelLimits(modelId, userId, tenantId);

    // 設定値とモデル制限の小さい方を選択
    const recommended = Math.min(
      currentBatchSize,
      limits.maxBatchSize,
      200, // 安全のための上限
    );

    if (recommended < currentBatchSize) {
      this.logger.warn(
        this.i18nService.formatMessage('batchSizeAdjusted', {
          old: currentBatchSize,
          new: recommended,
          limit: limits.maxBatchSize
        })
      );
    }

    return Math.max(10, recommended); // 最低10個
  }

  /**
   * チャンク数を推定
   */
  estimateChunkCount(textLength: number, chunkSize: number): number {
    const chunkSizeInChars = chunkSize * 4; // 1 token ≈ 4 chars
    return Math.ceil(textLength / chunkSizeInChars);
  }

  /**
   * ベクトル次元の検証
   */
  async validateDimensions(
    modelId: string,
    userId: string,
    actualDimensions: number,
    tenantId?: string,
  ): Promise<boolean> {
    const limits = await this.getModelLimits(modelId, userId, tenantId);

    if (actualDimensions !== limits.expectedDimensions) {
      this.logger.warn(
        this.i18nService.formatMessage('dimensionMismatch', {
          id: modelId,
          expected: limits.expectedDimensions,
          actual: actualDimensions
        })
      );
      return false;
    }

    return true;
  }

  /**
   * 設定概要を取得（ログ用）
   */
  async getConfigSummary(
    chunkSize: number,
    chunkOverlap: number,
    modelId: string,
    userId: string,
    tenantId?: string,
  ): Promise<string> {
    const limits = await this.getModelLimits(modelId, userId, tenantId);

    return [
      `モデル: ${modelId}`,
      `チャンクサイズ: ${chunkSize} tokens (制限: ${limits.maxInputTokens})`,
      `重なりサイズ: ${chunkOverlap} tokens`,
      `バッチサイズ: ${limits.maxBatchSize}`,
      `ベクトル次元: ${limits.expectedDimensions}`,
    ].join(', ');
  }

  /**
   * フロントエンド用の設定制限を取得
   * フロントエンドのスライダーの上限設定に使用
   */
  async getFrontendLimits(
    modelId: string,
    userId: string,
    tenantId?: string,
  ): Promise<{
    maxChunkSize: number;
    maxOverlapSize: number;
    minOverlapSize: number;
    defaultChunkSize: number;
    defaultOverlapSize: number;
    modelInfo: {
      name: string;
      maxInputTokens: number;
      maxBatchSize: number;
      expectedDimensions: number;
    };
  }> {
    const limits = await this.getModelLimits(modelId, userId, tenantId);

    // 最終的な上限を計算（環境変数とモデル制限の小さい方を選択）
    const maxChunkSize = Math.min(this.envMaxChunkSize, limits.maxInputTokens);
    const maxOverlapSize = Math.min(
      this.envMaxOverlapSize,
      Math.floor(maxChunkSize * this.DEFAULTS.maxOverlapRatio),
    );

    // モデル設定名を取得
    const modelConfig = await this.modelConfigService.findOne(modelId, userId, tenantId || '');
    const modelName = modelConfig?.name || 'Unknown';

    // テナントまたはユーザー設定からデフォルト値を取得
    let defaultChunkSize = this.DEFAULTS.chunkSize;
    let defaultOverlapSize = this.DEFAULTS.chunkOverlap;

    if (tenantId) {
      const tenantSettings = await this.tenantService.getSettings(tenantId);
      if (tenantSettings.chunkSize) defaultChunkSize = tenantSettings.chunkSize;
      if (tenantSettings.chunkOverlap) defaultOverlapSize = tenantSettings.chunkOverlap;
    }

    return {
      maxChunkSize,
      maxOverlapSize,
      minOverlapSize: this.DEFAULTS.minChunkOverlap,
      defaultChunkSize: Math.min(defaultChunkSize, maxChunkSize),
      defaultOverlapSize: Math.max(this.DEFAULTS.minChunkOverlap, Math.min(defaultOverlapSize, maxOverlapSize)),
      modelInfo: {
        name: modelName,
        maxInputTokens: limits.maxInputTokens,
        maxBatchSize: limits.maxBatchSize,
        expectedDimensions: limits.expectedDimensions,
      },
    };
  }
}