Yi.Admin/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Domain/Managers/UsageStatisticsManager.cs

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