diff --git a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/AIManagementOptions.cs b/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/AIManagementOptions.cs index 493fcb14a..72ab61ffd 100644 --- a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/AIManagementOptions.cs +++ b/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; } diff --git a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/ITokenUsageRecordRepository.cs b/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/ITokenUsageRecordRepository.cs new file mode 100644 index 000000000..11724dda4 --- /dev/null +++ b/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 +{ + Task FindByMessageIdAsync( + Guid conversationId, + Guid? messageId, + CancellationToken cancellationToken = default); + + Task GetCountAsync( + ISpecification specification, + CancellationToken cancellationToken = default); + + Task> GetListAsync( + ISpecification specification, + string sorting = $"{nameof(TokenUsageRecord.CreationTime)}", + int maxResultCount = 10, + int skipCount = 0, + CancellationToken cancellationToken = default); +} diff --git a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/TokenUsageRecord.cs b/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/TokenUsageRecord.cs index fedf9a194..136747a39 100644 --- a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/TokenUsageRecord.cs +++ b/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, 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; + } } diff --git a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/TokenUsageStore.cs b/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Tokens/TokenUsageStore.cs new file mode 100644 index 000000000..026ba82cd --- /dev/null +++ b/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); + } + } +} diff --git a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AIManagementDbContext.cs b/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AIManagementDbContext.cs index b43c95e49..fb3b6a9f2 100644 --- a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AIManagementDbContext.cs +++ b/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, IAIMan public DbSet WorkspaceDefinitions { get; set; } public DbSet TextChatMessageRecords { get; set; } public DbSet ConversationRecords { get; set; } + public DbSet TokenUsageRecords { get; set; } public AIManagementDbContext( DbContextOptions options) : base(options) { diff --git a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AIManagementDbContextModelBuilderExtensions.cs b/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AIManagementDbContextModelBuilderExtensions.cs index 79612bacf..b6e8d2110 100644 --- a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AIManagementDbContextModelBuilderExtensions.cs +++ b/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(b => + { + b.ToTable(AbpAIManagementDbProperties.DbTablePrefix + "TokenUsages", AbpAIManagementDbProperties.DbSchema); + + b.ConfigureByConvention(); + + b.HasIndex(x => new { x.TenantId, x.ConversationId }); + }); if (builder.IsHostDatabase()) { diff --git a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AbpAIManagementEntityFrameworkCoreModule.cs b/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AbpAIManagementEntityFrameworkCoreModule.cs index 0802cbc47..bb4cf3964 100644 --- a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/AbpAIManagementEntityFrameworkCoreModule.cs +++ b/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(); options.AddRepository(); + options.AddRepository(); options.AddRepository(); }); diff --git a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/EfCoreTokenUsageRecordRepository.cs b/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/EfCoreTokenUsageRecordRepository.cs new file mode 100644 index 000000000..2bb211cc6 --- /dev/null +++ b/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, ITokenUsageRecordRepository +{ + public EfCoreTokenUsageRecordRepository( + IDbContextProvider dbContextProvider) + : base(dbContextProvider) + { + } + + public async virtual Task 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 GetCountAsync( + ISpecification specification, + CancellationToken cancellationToken = default) + { + return await (await GetQueryableAsync()) + .Where(specification.ToExpression()) + .CountAsync(GetCancellationToken(cancellationToken)); + } + + public async virtual Task> GetListAsync( + ISpecification 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)); + } +} diff --git a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/IAIManagementDbContext.cs b/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/IAIManagementDbContext.cs index 33d66185c..edc46235d 100644 --- a/aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.EntityFrameworkCore/LINGYUN/Abp/AIManagement/EntityFrameworkCore/IAIManagementDbContext.cs +++ b/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 WorkspaceDefinitions { get; } DbSet TextChatMessageRecords { get; } DbSet ConversationRecords { get; } + DbSet TokenUsageRecords { get; } }