feat: Add persistent TokenUsage

aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/AIManagementOptions.cs

@@ -1,4 +1,6 @@
-namespace LINGYUN.Abp.AIManagement;
+using System;
+
+namespace LINGYUN.Abp.AIManagement;
 public class AIManagementOptions
 {
     public bool IsDynamicWorkspaceStoreEnabled {  get; set; }

aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/ITokenUsageRecordRepository.cs

@@ -0,0 +1,26 @@
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Volo.Abp.Domain.Repositories;
+using Volo.Abp.Specifications;
+
+namespace LINGYUN.Abp.AIManagement.Tokens;
+public interface ITokenUsageRecordRepository : IBasicRepository<TokenUsageRecord, Guid>
+{
+    Task<TokenUsageRecord?> FindByMessageIdAsync(
+        Guid conversationId,
+        Guid? messageId,
+        CancellationToken cancellationToken = default);
+
+    Task<int> GetCountAsync(
+        ISpecification<TokenUsageRecord> specification,
+        CancellationToken cancellationToken = default);
+
+    Task<List<TokenUsageRecord>> GetListAsync(
+        ISpecification<TokenUsageRecord> specification,
+        string sorting = $"{nameof(TokenUsageRecord.CreationTime)}",
+        int maxResultCount = 10,
+        int skipCount = 0,
+        CancellationToken cancellationToken = default);
+}

aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/TokenUsageRecord.cs

@@ -1,8 +1,42 @@
 using System;
-using System.Collections.Generic;
-using System.Text;
+using Volo.Abp.Domain.Entities.Auditing;
+using Volo.Abp.MultiTenancy;
 
 namespace LINGYUN.Abp.AIManagement.Tokens;
-public class TokenUsageRecord
+public class TokenUsageRecord : AuditedEntity<Guid>, IMultiTenant
 {
+    public Guid? TenantId { get; private set; }
+    public Guid? MessageId { get; private set; }
+    public Guid? ConversationId { get; private set; }
+    public long? InputTokenCount { get; set; }
+    public long? OutputTokenCount { get; set; }
+    public long? TotalTokenCount { get; set; }
+    public long? CachedInputTokenCount { get; set; }
+    public long? ReasoningTokenCount { get; set; }
+    protected TokenUsageRecord()
+    {
+
+    }
+
+    public TokenUsageRecord(
+        Guid id,
+        Guid? messageId = null,
+        Guid? conversationId = null, 
+        long? inputTokenCount = null,
+        long? outputTokenCount = null,
+        long? totalTokenCount = null,
+        long? cachedInputTokenCount = null,
+        long? reasoningTokenCount = null,
+        Guid? tenantId = null)
+        : base(id)
+    {
+        MessageId = messageId;
+        ConversationId = conversationId;
+        InputTokenCount = inputTokenCount;
+        OutputTokenCount = outputTokenCount;
+        TotalTokenCount = totalTokenCount;
+        CachedInputTokenCount = cachedInputTokenCount;
+        ReasoningTokenCount = reasoningTokenCount;
+        TenantId = tenantId;
+    }
 }

aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/TokenUsageStore.cs

@@ -0,0 +1,68 @@
+using LINGYUN.Abp.AI.Models;
+using LINGYUN.Abp.AI.Tokens;
+using System.Threading.Tasks;
+using Volo.Abp.DependencyInjection;
+using Volo.Abp.Guids;
+using Volo.Abp.MultiTenancy;
+
+namespace LINGYUN.Abp.AIManagement.Tokens;
+
+[Dependency(ReplaceServices = true)]
+public class TokenUsageStore : ITokenUsageStore, ITransientDependency
+{
+    private readonly ICurrentTenant _currentTenant;
+    private readonly IGuidGenerator _guidGenerator;
+    private readonly ITokenUsageRecordRepository _tokenUsageRecordRepository;
+
+    public TokenUsageStore(
+        ICurrentTenant currentTenant, 
+        IGuidGenerator guidGenerator, 
+        ITokenUsageRecordRepository tokenUsageRecordRepository)
+    {
+        _currentTenant = currentTenant;
+        _guidGenerator = guidGenerator;
+        _tokenUsageRecordRepository = tokenUsageRecordRepository;
+    }
+
+    public async virtual Task SaveTokenUsageAsync(TokenUsageInfo tokenUsageInfo)
+    {
+        var tokenUsageRecord = await _tokenUsageRecordRepository.FindByMessageIdAsync(
+            tokenUsageInfo.ConversationId,
+            tokenUsageInfo.MessageId);
+
+        if (tokenUsageRecord == null)
+        {
+            tokenUsageRecord = new TokenUsageRecord(
+                _guidGenerator.Create(),
+                tokenUsageInfo.MessageId,
+                tokenUsageInfo.ConversationId,
+                tokenUsageInfo.InputTokenCount,
+                tokenUsageInfo.OutputTokenCount,
+                tokenUsageInfo.TotalTokenCount,
+                tokenUsageInfo.CachedInputTokenCount,
+                tokenUsageInfo.ReasoningTokenCount,
+                _currentTenant.Id);
+
+            await _tokenUsageRecordRepository.InsertAsync(tokenUsageRecord);
+        }
+        else
+        {
+            tokenUsageRecord.InputTokenCount ??= 0;
+            tokenUsageRecord.InputTokenCount += tokenUsageInfo.InputTokenCount;
+
+            tokenUsageRecord.OutputTokenCount ??= 0;
+            tokenUsageRecord.OutputTokenCount += tokenUsageInfo.OutputTokenCount;
+
+            tokenUsageRecord.CachedInputTokenCount ??= 0;
+            tokenUsageRecord.CachedInputTokenCount += tokenUsageInfo.CachedInputTokenCount;
+
+            tokenUsageRecord.ReasoningTokenCount ??= 0;
+            tokenUsageRecord.ReasoningTokenCount += tokenUsageInfo.ReasoningTokenCount;
+
+            tokenUsageRecord.TotalTokenCount ??= 0;
+            tokenUsageRecord.TotalTokenCount += tokenUsageInfo.TotalTokenCount;
+
+            await _tokenUsageRecordRepository.UpdateAsync(tokenUsageRecord);
+        }
+    }
+}

aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AIManagementDbContext.cs

@@ -1,4 +1,5 @@
 using LINGYUN.Abp.AIManagement.Chats;
+using LINGYUN.Abp.AIManagement.Tokens;
 using LINGYUN.Abp.AIManagement.Workspaces;
 using Microsoft.EntityFrameworkCore;
 using Volo.Abp.Data;
@@ -12,6 +13,7 @@ public class AIManagementDbContext : AbpDbContext<AIManagementDbContext>, IAIMan
     public DbSet<WorkspaceDefinitionRecord> WorkspaceDefinitions { get; set; }
     public DbSet<TextChatMessageRecord> TextChatMessageRecords { get; set; }
     public DbSet<ConversationRecord> ConversationRecords { get; set; }
+    public DbSet<TokenUsageRecord> TokenUsageRecords { get; set; }
     public AIManagementDbContext(
         DbContextOptions<AIManagementDbContext> options) : base(options)
     {

aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AIManagementDbContextModelBuilderExtensions.cs

@@ -1,6 +1,7 @@
 using JetBrains.Annotations;
 using LINGYUN.Abp.AIManagement.Chats;
 using LINGYUN.Abp.AIManagement.EntityFrameworkCore.ValueConversions;
+using LINGYUN.Abp.AIManagement.Tokens;
 using LINGYUN.Abp.AIManagement.Workspaces;
 using Microsoft.EntityFrameworkCore;
 using Volo.Abp;
@@ -45,6 +46,14 @@ public static class AIManagementDbContextModelBuilderExtensions
 
             b.ApplyObjectExtensionMappings();
         });
+        builder.Entity<TokenUsageRecord>(b =>
+        {
+            b.ToTable(AbpAIManagementDbProperties.DbTablePrefix + "TokenUsages", AbpAIManagementDbProperties.DbSchema);
+
+            b.ConfigureByConvention();
+
+            b.HasIndex(x => new { x.TenantId, x.ConversationId });
+        });
 
         if (builder.IsHostDatabase())
         {

aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AbpAIManagementEntityFrameworkCoreModule.cs

@@ -1,4 +1,5 @@
 using LINGYUN.Abp.AIManagement.Chats;
+using LINGYUN.Abp.AIManagement.Tokens;
 using LINGYUN.Abp.AIManagement.Workspaces;
 using Microsoft.Extensions.DependencyInjection;
 using Volo.Abp.EntityFrameworkCore;
@@ -19,6 +20,7 @@ public class AbpAIManagementEntityFrameworkCoreModule : AbpModule
 
             options.AddRepository<ConversationRecord, EfCoreConversationRecordRepository>();
             options.AddRepository<TextChatMessageRecord, EfCoreTextChatMessageRecordRepository>();
+            options.AddRepository<TokenUsageRecord, EfCoreTokenUsageRecordRepository>();
 
             options.AddRepository<WorkspaceDefinitionRecord, EfCoreWorkspaceDefinitionRecordRepository>();
         });

aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/EfCoreTokenUsageRecordRepository.cs

@@ -0,0 +1,54 @@
+using LINGYUN.Abp.AIManagement.Tokens;
+using Microsoft.EntityFrameworkCore;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Linq.Dynamic.Core;
+using System.Threading;
+using System.Threading.Tasks;
+using Volo.Abp.Domain.Repositories.EntityFrameworkCore;
+using Volo.Abp.EntityFrameworkCore;
+using Volo.Abp.Specifications;
+
+namespace LINGYUN.Abp.AIManagement.EntityFrameworkCore;
+public class EfCoreTokenUsageRecordRepository : EfCoreRepository<IAIManagementDbContext, TokenUsageRecord, Guid>, ITokenUsageRecordRepository
+{
+    public EfCoreTokenUsageRecordRepository(
+        IDbContextProvider<IAIManagementDbContext> dbContextProvider) 
+        : base(dbContextProvider)
+    {
+    }
+
+    public async virtual Task<TokenUsageRecord?> FindByMessageIdAsync(
+        Guid conversationId,
+        Guid? messageId,
+        CancellationToken cancellationToken = default)
+    {
+        return await (await GetQueryableAsync())
+            .Where(x => x.ConversationId == conversationId && x.MessageId == messageId)
+            .FirstOrDefaultAsync(GetCancellationToken(cancellationToken));
+    }
+
+    public async virtual Task<int> GetCountAsync(
+        ISpecification<TokenUsageRecord> specification,
+        CancellationToken cancellationToken = default)
+    {
+        return await (await GetQueryableAsync())
+            .Where(specification.ToExpression())
+            .CountAsync(GetCancellationToken(cancellationToken));
+    }
+
+    public async virtual Task<List<TokenUsageRecord>> GetListAsync(
+        ISpecification<TokenUsageRecord> specification,
+        string sorting = $"{nameof(TokenUsageRecord.CreationTime)}",
+        int maxResultCount = 10,
+        int skipCount = 0,
+        CancellationToken cancellationToken = default)
+    {
+        return await (await GetQueryableAsync())
+            .Where(specification.ToExpression())
+            .OrderBy(!sorting.IsNullOrWhiteSpace() ? sorting : $"{nameof(TokenUsageRecord.CreationTime)}")
+            .PageBy(skipCount, maxResultCount)
+            .ToListAsync(GetCancellationToken(cancellationToken));
+    }
+}

aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/IAIManagementDbContext.cs

@@ -1,4 +1,5 @@
 using LINGYUN.Abp.AIManagement.Chats;
+using LINGYUN.Abp.AIManagement.Tokens;
 using LINGYUN.Abp.AIManagement.Workspaces;
 using Microsoft.EntityFrameworkCore;
 using Volo.Abp.Data;
@@ -12,4 +13,5 @@ public interface IAIManagementDbContext : IEfCoreDbContext
     DbSet<WorkspaceDefinitionRecord> WorkspaceDefinitions { get; }
     DbSet<TextChatMessageRecord> TextChatMessageRecords { get; }
     DbSet<ConversationRecord> ConversationRecords { get; }
+    DbSet<TokenUsageRecord> TokenUsageRecords { get; }
 }