aurak/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';
import { I18nService } from '../i18n/i18n.service';

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

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

  async findAll(): Promise<ModelConfig[]> {
    return this.modelConfigRepository.find();
  }

  async findOne(id: string): Promise<ModelConfig> {
    const modelConfig = await this.modelConfigRepository.findOne({
      where: { id },
    });

    if (!modelConfig) {
      throw new NotFoundException(
        this.i18nService.formatMessage('modelConfigNotFound', { id }),
      );
    }
    return modelConfig;
  }

  async findByType(type: string): Promise<ModelConfig[]> {
    return this.modelConfigRepository.find({ where: { type } });
  }

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

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

  async remove(id: string): Promise<void> {
    const result = await this.modelConfigRepository.delete({ id });
    if (result.affected === 0) {
      throw new NotFoundException(
        this.i18nService.formatMessage('modelConfigNotFound', { id }),
      );
    }
  }

  /**
   * Set the specified model as default
   */
  async setDefault(id: string): Promise<ModelConfig> {
    const modelConfig = await this.findOne(id);

    // Clear default flag for other models of the same type
    await this.modelConfigRepository
      .createQueryBuilder()
      .update(ModelConfig)
      .set({ isDefault: false })
      .where('type = :type', { type: modelConfig.type })
      .execute();

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

  /**
   * Get default model for specified type
   * Strict rule: Only return models specified in Index Chat Config, throw error if not found
   */
  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;
  }
}