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;
}
}