ai-center/src/provider/utils/parse-config.ts

import CryptoJS from 'crypto-js';

// 加密函数
export function encryptAES(plainText: string, secretKey: string) {
  return CryptoJS.AES.encrypt(plainText, secretKey).toString();
}

// 解密函数
export function decryptAES(cipherText: string, secretKey: string) {
  const bytes = CryptoJS.AES.decrypt(cipherText, secretKey);
  return bytes.toString(CryptoJS.enc.Utf8);
}

type AIModel = {
  /**
   * 提供商
   */
  provider: string;
  /**
   * 模型名称
   */
  model: string;
  /**
   * 模型组
   */
  group: string;
  /**
   * 每日请求频率限制
   */
  dayLimit?: number;
  /**
   * 总的token限制
   */
  tokenLimit?: number;
};

type SecretKey = {
  /**
   * 组
   */
  group: string;
  /**
   * API密钥
   */
  apiKey: string;
  /**
   * 解密密钥
   */
  decryptKey?: string;
};

export type GetProviderOpts = {
  model: string;
  group: string;
  decryptKey?: string;
};
export type ProviderResult = {
  provider: string;
  model: string;
  group: string;
  apiKey: string;
  dayLimit?: number;
  tokenLimit?: number;
  baseURL?: string;
  /**
   * 解密密钥
   */
  decryptKey?: string;
};

export type AIConfig = {
  title?: string;
  description?: string;
  models: AIModel[];
  secretKeys: SecretKey[];
  filter?: {
    objectKey: string;
    type: 'array' | 'object';
    operate: 'removeAttribute' | 'remove';
    attribute: string[];
  }[];
};
export class AIConfigParser {
  private config: AIConfig;
  result: ProviderResult;
  constructor(config: AIConfig) {
    this.config = config;
  }

  getProvider(opts: GetProviderOpts): ProviderResult {
    const { model, group, decryptKey } = opts;
    const modelConfig = this.config.models.find((m) => m.model === model && m.group === group);
    const groupConfig = this.config.secretKeys.find((m) => m.group === group);
    if (!modelConfig) {
      throw new Error(`在模型组 ${group} 中未找到模型 ${model}`);
    }
    const mergeConfig = {
      ...modelConfig,
      ...groupConfig,
      decryptKey: decryptKey || groupConfig?.decryptKey,
    };
    // 验证模型配置
    if (!mergeConfig.provider) {
      throw new Error(`模型 ${model} 未配置提供商`);
    }
    if (!mergeConfig.model) {
      throw new Error(`模型 ${model} 未配置模型名称`);
    }
    if (!mergeConfig.apiKey) {
      throw new Error(`组 ${group} 未配置 API 密钥`);
    }
    if (!mergeConfig.group) {
      throw new Error(`组 ${group} 未配置`);
    }
    this.result = mergeConfig;
    return mergeConfig;
  }

  async getSecretKey({
    getCache,
    setCache,
    providerResult,
  }: {
    getCache?: (key: string) => Promise<string>;
    setCache?: (key: string, value: string) => Promise<void>;
    providerResult?: ProviderResult;
  }) {
    const { apiKey, decryptKey, group = '', model } = providerResult || this.result;
    const cacheKey = `${group}--${model}`;
    if (!decryptKey) {
      return apiKey;
    }
    if (getCache) {
      const cache = await getCache(cacheKey);
      if (cache) {
        return cache;
      }
    }
    const secretKey = decryptAES(apiKey, decryptKey);
    if (setCache) {
      await setCache(cacheKey, secretKey);
    }
    return secretKey;
  }
  encrypt(plainText: string, secretKey: string) {
    return encryptAES(plainText, secretKey);
  }

  decrypt(cipherText: string, secretKey: string) {
    return decryptAES(cipherText, secretKey);
  }
}