ai-center/src/routes/ai-chat/services/chat-config-srevices.ts

import { AIConfigParser, type AIConfig } from '@/provider/utils/parse-config.ts';
import { redis } from '@/modules/db.ts';
import { CustomError } from '@kevisual/router';
import { queryConfig } from '@/modules/query.ts';
import { log } from '@/logger/index.ts';
export class ChatConfigServices {
  cachePrefix = 'ai:chat:config';
  // 使用谁的模型
  owner: string;
  // 使用者
  username: string;
  aiConfig?: AIConfig;
  isOwner: boolean;
  /**
   * username 是使用的模型的用户名，使用谁的模型
   * @param username
   */
  constructor(owner: string, username: string, token?: string) {
    this.owner = owner;
    this.username = username;
    this.isOwner = owner === username;
  }
  getKey() {
    return `${this.cachePrefix}:${this.owner}`;
  }
  /**
   * 获取chat配置
   * @param keepSecret 是否需要清除secret 默认 不清除 为true
   * @returns
   */
  async getChatConfig(keepSecret = true, token?: string) {
    const key = this.getKey();
    const cache = await redis.get(key);
    let modelConfig = null;
    if (cache) {
      modelConfig = JSON.parse(cache);
    }
    if (!modelConfig) {
      if (this.owner !== this.username) {
        throw new CustomError(
          `the owner [${this.owner}] config, [${this.username}] not permission to init config, only owner can init config, place connect owner`,
        );
      } else {
        const res = await queryConfig.getConfigByKey('ai.json', { token });
        if (res.code === 200 && res.data?.data) {
          modelConfig = res.data.data;
        } else {
          throw new CustomError(400, 'get config failed');
        }
      }
    }
    if (!modelConfig) {
      throw new CustomError(`${this.owner} modelConfig is null`);
    }
    if (!cache) {
      const cacheTime = 60 * 60 * 24 * 40; // 1天
      await redis.set(key, JSON.stringify(modelConfig), 'EX', cacheTime);
    }
    this.aiConfig = modelConfig;
    if (!keepSecret) {
      modelConfig = this.filterApiKey(modelConfig);
    }
    return modelConfig;
  }
  async clearCache() {
    const key = this.getKey();
    await redis.set(key, JSON.stringify({}), 'EX', 1);
  }
  /**
   * 获取模型配置
   * @returns
   */
  async getSelectOpts(config?: AIConfig) {
    const aiConfigParser = new AIConfigParser(config || this.aiConfig);
    return aiConfigParser.getSelectOpts();
  }
  async filterApiKey(chatConfig: AIConfig) {
    // 过滤掉secret中的所有apiKey，移除掉并返回chatConfig
    const { secretKeys = [], ...rest } = chatConfig;
    return {
      ...rest,
      secretKeys: secretKeys.map((item) => {
        return {
          ...item,
          apiKey: undefined,
          decryptKey: undefined,
        };
      }),
    };
  }
  /**
   * 获取和检测当前用户的额度, 当使用 root 账号的时候，才需要检测
   * username是当前使用用户
   * @param username
   */
  async checkUserCanChat(username: string) {
    if (this.owner !== 'root') return true;
    const maxToken = 100000;
    const userCacheKey = `${this.cachePrefix}:root:chat-limit:${username}`;
    const cache = await redis.get(userCacheKey);
    if (cache) {
      const cacheData = JSON.parse(cache);
      if (cacheData.token >= maxToken) {
        throw new CustomError(400, 'use root account token limit exceeded');
      }
    }
    return true;
  }
  /**
   * 获取用户的使用情况
   * username是当前使用用户
   * @param username
   * @returns
   */
  async getUserChatLimit(username: string) {
    if (this.owner !== 'root') return;
    const userCacheKey = `${this.cachePrefix}:root:chat-limit:${username}`;
    const cache = await redis.get(userCacheKey);
    if (cache) {
      const cacheData = JSON.parse(cache);
      return cacheData;
    }
    return {
      token: 0,
      day: 0,
    };
  }

  /**
   * 更新用户的使用情况
   * username是当前使用用户
   * @param username
   * @param token
   */
  async updateUserChatLimit(username: string, token: number) {
    if (this.owner !== 'root') return;
    const userCacheKey = `${this.cachePrefix}:root:chat-limit:${username}`;
    const cache = await redis.get(userCacheKey);
    if (cache) {
      const cacheData = JSON.parse(cache);
      cacheData.token = cacheData.token + token;
      await redis.set(userCacheKey, JSON.stringify(cacheData), 'EX', 60 * 60 * 24 * 30); // 30天
    } else {
      await redis.set(userCacheKey, JSON.stringify({ token }), 'EX', 60 * 60 * 24 * 30); // 30天
    }
  }
  async clearChatLimit() {
    if (this.owner !== 'root') return;
    // const userCacheKey = `${this.cachePrefix}:root:chat-limit:${this.username}`;
    // await redis.del(userCacheKey);
  }
}