using SqlSugar; using Volo.Abp.Domain.Services; using Yi.Framework.AiHub.Domain.Entities; using Yi.Framework.AiHub.Domain.Entities.Model; using Yi.Framework.AiHub.Domain.Entities.OpenApi; using Yi.Framework.AiHub.Domain.Shared.Consts; using Yi.Framework.SqlSugarCore.Abstractions; namespace Yi.Framework.AiHub.Domain.Managers; /// /// Token验证结果 /// public class TokenValidationResult { /// /// 用户Id /// public Guid UserId { get; set; } /// /// Token Id /// public Guid TokenId { get; set; } /// /// token /// public string Token { get; set; } } public class TokenManager : DomainService { private readonly ISqlSugarRepository _tokenRepository; private readonly ISqlSugarRepository _usageStatisticsRepository; private readonly ISqlSugarRepository _aiModelRepository; public TokenManager( ISqlSugarRepository tokenRepository, ISqlSugarRepository usageStatisticsRepository, ISqlSugarRepository aiModelRepository) { _tokenRepository = tokenRepository; _usageStatisticsRepository = usageStatisticsRepository; _aiModelRepository = aiModelRepository; } /// /// 验证Token并返回用户Id和TokenId /// /// Token密钥或者TokenId /// 模型Id(用于判断是否是尊享模型需要检查额度) /// Token验证结果 public async Task ValidateTokenAsync(object tokenOrId, string? modelId = null) { if (tokenOrId is null) { throw new UserFriendlyException("当前请求未包含token", "401"); } TokenAggregateRoot entity; if (tokenOrId is Guid tokenId) { entity = await _tokenRepository._DbQueryable .Where(x => x.Id == tokenId) .FirstAsync(); } else { var tokenStr = tokenOrId.ToString(); if (!tokenStr.StartsWith("yi-")) { throw new UserFriendlyException("当前请求token非法", "401"); } entity = await _tokenRepository._DbQueryable .Where(x => x.Token == tokenStr) .FirstAsync(); } if (entity is null) { throw new UserFriendlyException("当前请求token无效", "401"); } // 检查Token是否被禁用 if (entity.IsDisabled) { throw new UserFriendlyException("当前Token已被禁用,请启用后再使用", "403"); } // 检查Token是否过期 if (entity.ExpireTime.HasValue && entity.ExpireTime.Value < DateTime.Now) { throw new UserFriendlyException("当前Token已过期,请更新过期时间或创建新的Token", "403"); } // 如果是尊享模型且Token设置了额度限制,检查是否超限 if (!string.IsNullOrEmpty(modelId) && entity.PremiumQuotaLimit.HasValue) { var isPremium = await _aiModelRepository._DbQueryable .Where(x => x.ModelId == modelId) .Select(x => x.IsPremium) .FirstAsync(); if (isPremium) { var usedQuota = await GetTokenPremiumUsedQuotaAsync(entity.UserId, entity.Id); if (usedQuota >= entity.PremiumQuotaLimit.Value) { throw new UserFriendlyException($"当前Token的尊享包额度已用完(已使用:{usedQuota},限制:{entity.PremiumQuotaLimit.Value}),请调整额度限制或使用其他Token", "403"); } } } return new TokenValidationResult { UserId = entity.UserId, TokenId = entity.Id, Token = entity.Token }; } /// /// 获取Token的尊享包已使用额度 /// private async Task GetTokenPremiumUsedQuotaAsync(Guid userId, Guid tokenId) { // 先获取所有尊享模型的ModelId列表 var premiumModelIds = await _aiModelRepository._DbQueryable .Where(x => x.IsPremium) .Select(x => x.ModelId) .ToListAsync(); var usedQuota = await _usageStatisticsRepository._DbQueryable .Where(x => x.UserId == userId && x.TokenId == tokenId && premiumModelIds.Contains(x.ModelId)) .SumAsync(x => x.TotalTokenCount); return usedQuota; } /// /// 获取用户的Token(兼容旧接口,返回第一个可用的Token) /// [Obsolete("请使用 ValidateTokenAsync 方法")] public async Task GetAsync(Guid userId) { var entity = await _tokenRepository._DbQueryable .Where(x => x.UserId == userId && !x.IsDisabled) .OrderBy(x => x.CreationTime) .FirstAsync(); return entity?.Token; } /// /// 获取用户Id(兼容旧接口) /// [Obsolete("请使用 ValidateTokenAsync 方法")] public async Task GetUserIdAsync(string? token) { var result = await ValidateTokenAsync(token); return result.UserId; } }