KnowledgeBase/server/src/model-config/model-config.service.ts

import { Injectable, NotFoundException, ForbiddenException, BadRequestException, forwardRef, Inject } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { Repository } from 'typeorm';
import { ModelConfig } from './model-config.entity';
import { CreateModelConfigDto } from './dto/create-model-config.dto';
import { UpdateModelConfigDto } from './dto/update-model-config.dto';
import { GLOBAL_TENANT_ID } from '../common/constants';
import { TenantService } from '../tenant/tenant.service';
import { ModelType } from '../types';

@Injectable()
export class ModelConfigService {
  constructor(
    @InjectRepository(ModelConfig)
    private modelConfigRepository: Repository<ModelConfig>,
    @Inject(forwardRef(() => TenantService))
    private readonly tenantService: TenantService,
  ) { }

  async create(
    userId: string,
    tenantId: string,
    createModelConfigDto: CreateModelConfigDto,
  ): Promise<ModelConfig> {
    const modelConfig = this.modelConfigRepository.create({
      ...createModelConfigDto,
      userId,
      tenantId,
    });
    return this.modelConfigRepository.save(modelConfig);
  }

  async findAll(userId: string, tenantId: string): Promise<ModelConfig[]> {
    return this.modelConfigRepository.createQueryBuilder('model')
      .where('model.tenantId = :tenantId OR model.tenantId IS NULL OR model.tenantId = :globalTenantId', {
        tenantId,
        globalTenantId: GLOBAL_TENANT_ID
      })
      .getMany();
  }

  async findOne(id: string, userId: string, tenantId: string): Promise<ModelConfig> {
    const modelConfig = await this.modelConfigRepository.createQueryBuilder('model')
      .where('model.id = :id', { id })
      .andWhere('(model.tenantId = :tenantId OR model.tenantId IS NULL OR model.tenantId = :globalTenantId)', {
        tenantId,
        globalTenantId: GLOBAL_TENANT_ID
      })
      .getOne();

    if (!modelConfig) {
      throw new NotFoundException(
        `ModelConfig with ID "${id}" not found.`,
      );
    }
    return modelConfig;
  }

  async findByType(userId: string, tenantId: string, type: string): Promise<ModelConfig[]> {
    return this.modelConfigRepository.createQueryBuilder('model')
      .where('model.type = :type', { type })
      .andWhere('(model.tenantId = :tenantId OR model.tenantId IS NULL OR model.tenantId = :globalTenantId)', {
        tenantId,
        globalTenantId: GLOBAL_TENANT_ID
      })
      .getMany();
  }

  async update(
    userId: string,
    tenantId: string,
    id: string,
    updateModelConfigDto: UpdateModelConfigDto,
  ): Promise<ModelConfig> {
    const modelConfig = await this.findOne(id, userId, tenantId);

    if (!modelConfig) {
      throw new NotFoundException(
        `ModelConfig with ID "${id}" not found.`,
      );
    }

    // Only allow updating if it belongs to the tenant, or if it's a global admin (not fully implemented, so we check tenantId)
    if (modelConfig.tenantId && modelConfig.tenantId !== tenantId) {
      throw new ForbiddenException('Cannot update models from another tenant');
    }

    // Update the model
    const updated = this.modelConfigRepository.merge(
      modelConfig,
      updateModelConfigDto,
    );
    return this.modelConfigRepository.save(updated);
  }

  async remove(userId: string, tenantId: string, id: string): Promise<void> {
    // Only allow removing if it exists and accessible in current tenant context
    const model = await this.findOne(id, userId, tenantId);
    if (model.tenantId && model.tenantId !== tenantId) {
      throw new ForbiddenException('Cannot delete models from another tenant');
    }
    const result = await this.modelConfigRepository.delete({ id });
    if (result.affected === 0) {
      throw new NotFoundException(`ModelConfig with ID "${id}" not found.`);
    }
  }

  /**
   * 指定されたモデルをデフォルトに設定
   */
  async setDefault(userId: string, tenantId: string, id: string): Promise<ModelConfig> {
    const modelConfig = await this.findOne(id, userId, tenantId);

    // 同じタイプの他のモデルのデフォルトフラグをクリア (現在のテナント内またはglobal)
    // 厳密には、現在のテナントのIsDefault設定といった方が正しいですが、シンプルにするため全体のIsDefaultを操作します
    await this.modelConfigRepository
      .createQueryBuilder()
      .update(ModelConfig)
      .set({ isDefault: false })
      .where('type = :type', { type: modelConfig.type })
      .andWhere('(tenantId = :tenantId OR tenantId IS NULL OR tenantId = :globalTenantId)', {
        tenantId,
        globalTenantId: GLOBAL_TENANT_ID
      })
      .execute();

    modelConfig.isDefault = true;
    return this.modelConfigRepository.save(modelConfig);
  }

  /**
   * 指定されたタイプのデフォルトモデルを取得
   * 厳密なルール：Index Chat Configで指定されたモデルのみを返し、なければエラーを投げる
   */
  async findDefaultByType(tenantId: string, type: ModelType): Promise<ModelConfig> {
    const settings = await this.tenantService.getSettings(tenantId);
    if (!settings) {
      throw new BadRequestException(`Organization settings not found for tenant: ${tenantId}`);
    }

    let modelId: string | undefined;
    if (type === ModelType.LLM) {
      modelId = settings.selectedLLMId;
    } else if (type === ModelType.EMBEDDING) {
      modelId = settings.selectedEmbeddingId;
    } else if (type === ModelType.RERANK) {
      modelId = settings.selectedRerankId;
    }

    if (!modelId) {
      throw new BadRequestException(`Model of type "${type}" is not configured in Index Chat Config for this organization.`);
    }

    const model = await this.modelConfigRepository.findOne({
      where: { id: modelId, isEnabled: true }
    });

    if (!model) {
      throw new BadRequestException(`The configured model for "${type}" (ID: ${modelId}) is either missing or disabled in model management.`);
    }

    return model;
  }
}