feat: 兼容了用量使用显示

This commit is contained in:
chenchun
2025-07-22 10:40:23 +08:00
parent b6d670c240
commit a0eb234539
8 changed files with 97 additions and 75 deletions

View File

@@ -132,8 +132,7 @@ public class AiGateWayManager : DomainService
TokenUsage = data.Usage
});
await _usageStatisticsManager.SetUsageAsync(userId.Value, request.Model, data.Usage.InputTokens ?? 0,
data.Usage.OutputTokens ?? 0);
await _usageStatisticsManager.SetUsageAsync(userId.Value, request.Model, data.Usage);
}
await response.WriteAsJsonAsync(data, cancellationToken);
@@ -200,7 +199,7 @@ public class AiGateWayManager : DomainService
{
await foreach (var data in completeChatResponse)
{
if (data.Usage is not null && data.Usage.TotalTokens is not null)
if (data.Usage is not null)
{
tokenUsage = data.Usage;
}
@@ -261,8 +260,7 @@ public class AiGateWayManager : DomainService
TokenUsage = tokenUsage
});
await _usageStatisticsManager.SetUsageAsync(userId.Value, request.Model, tokenUsage.InputTokens ?? 0,
tokenUsage.OutputTokens ?? 0);
await _usageStatisticsManager.SetUsageAsync(userId.Value, request.Model, tokenUsage);
}
}
}

View File

@@ -1,5 +1,6 @@
using Medallion.Threading;
using Volo.Abp.Domain.Services;
using Yi.Framework.AiHub.Application.Contracts.Dtos.OpenAi;
using Yi.Framework.AiHub.Domain.Entities;
using Yi.Framework.SqlSugarCore.Abstractions;
@@ -17,8 +18,16 @@ public class UsageStatisticsManager : DomainService
private IDistributedLockProvider DistributedLock =>
LazyServiceProvider.LazyGetRequiredService<IDistributedLockProvider>();
public async Task SetUsageAsync(Guid userId, string modelId, int inputTokenCount, int outputTokenCount)
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);
@@ -37,8 +46,4 @@ public class UsageStatisticsManager : DomainService
}
}
}
}
internal class LazyServiceProvider
{
}