49 lines
1.8 KiB
C#
49 lines
1.8 KiB
C#
using Medallion.Threading;
|
|
using Volo.Abp.Domain.Services;
|
|
using Yi.Framework.AiHub.Domain.Entities;
|
|
using Yi.Framework.AiHub.Domain.Shared.Dtos.OpenAi;
|
|
using Yi.Framework.SqlSugarCore.Abstractions;
|
|
|
|
namespace Yi.Framework.AiHub.Domain.Managers;
|
|
|
|
public class UsageStatisticsManager : DomainService
|
|
{
|
|
private readonly ISqlSugarRepository<UsageStatisticsAggregateRoot> _repository;
|
|
|
|
public UsageStatisticsManager(ISqlSugarRepository<UsageStatisticsAggregateRoot> repository)
|
|
{
|
|
_repository = repository;
|
|
}
|
|
|
|
private IDistributedLockProvider DistributedLock =>
|
|
LazyServiceProvider.LazyGetRequiredService<IDistributedLockProvider>();
|
|
|
|
public async Task SetUsageAsync(Guid? userId, string modelId, ThorUsageResponse? tokenUsage)
|
|
{
|
|
long inputTokenCount = tokenUsage?.PromptTokens
|
|
?? tokenUsage?.InputTokens
|
|
?? 0;
|
|
|
|
long outputTokenCount = tokenUsage?.CompletionTokens
|
|
?? tokenUsage?.OutputTokens
|
|
?? 0;
|
|
|
|
await using (await DistributedLock.AcquireLockAsync($"UsageStatistics:{userId?.ToString()}"))
|
|
{
|
|
var entity = await _repository._DbQueryable.FirstAsync(x => x.UserId == userId && x.ModelId == modelId);
|
|
//存在数据，更细
|
|
if (entity is not null)
|
|
{
|
|
entity.AddOnceChat(inputTokenCount, outputTokenCount);
|
|
await _repository.UpdateAsync(entity);
|
|
}
|
|
//不存在插入
|
|
else
|
|
{
|
|
var usage = new UsageStatisticsAggregateRoot(userId, modelId);
|
|
usage.AddOnceChat(inputTokenCount, outputTokenCount);
|
|
await _repository.InsertAsync(usage);
|
|
}
|
|
}
|
|
}
|
|
} |