fix: 修复用量统计线程问题并完善搜索与Token计算逻辑

- OnlineSearch 增加 daysAgo 非法值保护，避免无效时间范围 - 修复 UsageStatistics 中 Prompt/Completion Token 为 0 时的统计异常 - 引入独立 UnitOfWork，解决流式处理下的并发与事务问题 - 确保用量统计、系统消息与尊享包扣减的原子性 - 补充前端 Element Plus 组件类型声明 - 统一并优化部分代码格式，不影响业务逻辑
2026-01-08 23:46:57 +08:00
parent 2f1f25ca37
commit 2544c01e9d
4 changed files with 92 additions and 86 deletions
--- a/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Domain/Managers/ChatManager.cs
+++ b/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Domain/Managers/ChatManager.cs
@@ -8,17 +8,15 @@ using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
-using ModelContextProtocol.Server;
 using OpenAI;
 using OpenAI.Chat;
 using Volo.Abp.Domain.Services;
+using Volo.Abp.Uow;
 using Yi.Framework.AiHub.Application.Contracts.Dtos.Chat;
 using Yi.Framework.AiHub.Domain.AiGateWay;
 using Yi.Framework.AiHub.Domain.Entities.Chat;
 using Yi.Framework.AiHub.Domain.Entities.Model;
-using Yi.Framework.AiHub.Domain.Entities.OpenApi;
 using Yi.Framework.AiHub.Domain.Shared.Attributes;
-using Yi.Framework.AiHub.Domain.Shared.Consts;
 using Yi.Framework.AiHub.Domain.Shared.Dtos;
 using Yi.Framework.AiHub.Domain.Shared.Dtos.OpenAi;
 using Yi.Framework.AiHub.Domain.Shared.Enums;
@@ -36,12 +34,14 @@ public class ChatManager : DomainService
    private readonly PremiumPackageManager _premiumPackageManager;
    private readonly AiGateWayManager _aiGateWayManager;
    private readonly ISqlSugarRepository<AiModelEntity, Guid> _aiModelRepository;
+    private readonly IUnitOfWorkManager _unitOfWorkManager;

    public ChatManager(ILoggerFactory loggerFactory,
        ISqlSugarRepository<MessageAggregateRoot> messageRepository,
        ISqlSugarRepository<AgentStoreAggregateRoot> agentStoreRepository, AiMessageManager aiMessageManager,
        UsageStatisticsManager usageStatisticsManager, PremiumPackageManager premiumPackageManager,
-        AiGateWayManager aiGateWayManager, ISqlSugarRepository<AiModelEntity, Guid> aiModelRepository)
+        AiGateWayManager aiGateWayManager, ISqlSugarRepository<AiModelEntity, Guid> aiModelRepository,
+        IUnitOfWorkManager unitOfWorkManager)
    {
        _loggerFactory = loggerFactory;
        _messageRepository = messageRepository;
@@ -51,6 +51,7 @@ public class ChatManager : DomainService
        _premiumPackageManager = premiumPackageManager;
        _aiGateWayManager = aiGateWayManager;
        _aiModelRepository = aiModelRepository;
+        _unitOfWorkManager = unitOfWorkManager;
    }

    /// <summary>
@@ -189,53 +190,59 @@ public class ChatManager : DomainService

                    //用量统计
                    case UsageContent usageContent:
-                        var usage = new ThorUsageResponse
+                        //由于MAF线程问题
+                        using (var uow = _unitOfWorkManager.Begin(requiresNew: true))
                        {
-                            InputTokens = Convert.ToInt32(usageContent.Details.InputTokenCount ?? 0),
-                            OutputTokens = Convert.ToInt32(usageContent.Details.OutputTokenCount ?? 0),
-                            TotalTokens = usageContent.Details.TotalTokenCount ?? 0,
-                        };
-                        //设置倍率
-                        usage.SetSupplementalMultiplier(modelDescribe.Multiplier);
-
-                        //创建系统回答，用于计费统计
-                        await _aiMessageManager.CreateSystemMessageAsync(userId, sessionId, new MessageInputDto
-                        {
-                            Content = "不与存储",
-                            ModelId = modelId,
-                            TokenUsage = usage
-                        }, tokenId);
-
-                        //创建用量统计，用于统计分析
-                        await _usageStatisticsManager.SetUsageAsync(userId, modelId, usage, tokenId);
-
-                        //扣减尊享token包用量
-                        var isPremium = await _aiModelRepository._DbQueryable
-                            .Where(x => x.ModelId == modelId)
-                            .Select(x => x.IsPremium)
-                            .FirstAsync();
-
-                        if (isPremium)
-                        {
-                            var totalTokens = usage?.TotalTokens ?? 0;
-                            if (totalTokens > 0)
+                            var usage = new ThorUsageResponse
                            {
-                                await _premiumPackageManager.TryConsumeTokensAsync(userId, totalTokens);
+                                InputTokens = Convert.ToInt32(usageContent.Details.InputTokenCount ?? 0),
+                                OutputTokens = Convert.ToInt32(usageContent.Details.OutputTokenCount ?? 0),
+                                TotalTokens = usageContent.Details.TotalTokenCount ?? 0,
+                            };
+                            //设置倍率
+                            usage.SetSupplementalMultiplier(modelDescribe.Multiplier);
+
+                            //创建系统回答，用于计费统计
+                            await _aiMessageManager.CreateSystemMessageAsync(userId, sessionId, new MessageInputDto
+                            {
+                                Content = "不与存储",
+                                ModelId = modelId,
+                                TokenUsage = usage
+                            }, tokenId);
+
+                            //创建用量统计，用于统计分析
+                            await _usageStatisticsManager.SetUsageAsync(userId, modelId, usage, tokenId);
+
+                            //扣减尊享token包用量
+                            var isPremium = await _aiModelRepository._DbQueryable
+                                .Where(x => x.ModelId == modelId)
+                                .Select(x => x.IsPremium)
+                                .FirstAsync();
+
+                            if (isPremium)
+                            {
+                                var totalTokens = usage?.TotalTokens ?? 0;
+                                if (totalTokens > 0)
+                                {
+                                    await _premiumPackageManager.TryConsumeTokensAsync(userId, totalTokens);
+                                }
                            }
-                        }

-                        await SendHttpStreamMessageAsync(httpContext,
-                            new AgentResultOutput
-                            {
-                                TypeEnum = update.RawRepresentation is ChatResponseUpdate raw
-                                    ? raw.FinishReason?.Value == "tool_calls"
-                                        ? AgentResultTypeEnum.ToolCallUsage
-                                        : AgentResultTypeEnum.Usage
-                                    : AgentResultTypeEnum.Usage,
-                                Content = usage!
-                            },
-                            isDone: false, cancellationToken);
-                        break;
+                            await uow.CompleteAsync();
+
+                            await SendHttpStreamMessageAsync(httpContext,
+                                new AgentResultOutput
+                                {
+                                    TypeEnum = update.RawRepresentation is ChatResponseUpdate raw
+                                        ? raw.FinishReason?.Value == "tool_calls"
+                                            ? AgentResultTypeEnum.ToolCallUsage
+                                            : AgentResultTypeEnum.Usage
+                                        : AgentResultTypeEnum.Usage,
+                                    Content = usage!
+                                },
+                                isDone: false, cancellationToken);
+                            break;
+                        }
                }
            }
        }
@@ -247,8 +254,13 @@ public class ChatManager : DomainService
        string serializedJson = currentThread.Serialize(JsonSerializerOptions.Web).GetRawText();
        agentStore.Store = serializedJson;

-        //插入或者更新
-        await _agentStoreRepository.InsertOrUpdateAsync(agentStore);
+        //由于MAF线程问题
+        using (var uow = _unitOfWorkManager.Begin(requiresNew: true))
+        {
+            //插入或者更新
+            await _agentStoreRepository.InsertOrUpdateAsync(agentStore);
+            await uow.CompleteAsync();
+        }
    }


--- a/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Domain/Managers/UsageStatisticsManager.cs
+++ b/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Domain/Managers/UsageStatisticsManager.cs
@@ -22,13 +22,13 @@ public class UsageStatisticsManager : DomainService
    {
        var actualTokenId = tokenId ?? Guid.Empty;

-        long inputTokenCount = tokenUsage?.PromptTokens
-                               ?? tokenUsage?.InputTokens
-                               ?? 0;
+        long inputTokenCount = tokenUsage?.PromptTokens > 0
+            ? tokenUsage.PromptTokens.Value
+            : tokenUsage?.InputTokens ?? 0;

-        long outputTokenCount = tokenUsage?.CompletionTokens
-                                ?? tokenUsage?.OutputTokens
-                                ?? 0;
+        long outputTokenCount = tokenUsage?.CompletionTokens > 0
+            ? tokenUsage.CompletionTokens.Value
+            : tokenUsage?.OutputTokens ?? 0;

        await using (await DistributedLock.AcquireLockAsync($"UsageStatistics:{userId?.ToString()}:{actualTokenId}:{modelId}"))
        {
--- a/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Domain/Mcp/OnlineSearchTool.cs
+++ b/Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Domain/Mcp/OnlineSearchTool.cs
@@ -29,8 +29,14 @@ public class OnlineSearchTool : ISingletonDependency
    }

    [YiAgentTool("联网搜索"), DisplayName("OnlineSearch"), Description("进行在线搜索，获取最新的网络信息，近期信息是7天，实时信息是1天")]
-    public async Task<string> OnlineSearch([Description("搜索关键字")]string keyword, [Description("距离现在多久天")]int? daysAgo = null)
+    public async Task<string> OnlineSearch([Description("搜索关键字")] string keyword,
+        [Description("距离现在多久天")] int? daysAgo = null)
    {
+        if (daysAgo <= 0)
+        {
+            daysAgo = 1;
+        }
+
        if (string.IsNullOrWhiteSpace(keyword))
        {
            return "搜索关键词不能为空";
@@ -149,8 +155,7 @@ public class OnlineSearchTool : ISingletonDependency
 /// </summary>
 public class BaiduSearchRequest
 {
-    [JsonPropertyName("messages")]
-    public List<BaiduSearchMessage> Messages { get; set; } = new();
+    [JsonPropertyName("messages")] public List<BaiduSearchMessage> Messages { get; set; } = new();

    [JsonPropertyName("search_filter")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
@@ -162,8 +167,7 @@ public class BaiduSearchRequest
 /// </summary>
 public class BaiduSearchFilter
 {
-    [JsonPropertyName("range")]
-    public BaiduSearchRange? Range { get; set; }
+    [JsonPropertyName("range")] public BaiduSearchRange? Range { get; set; }
 }

 /// <summary>
@@ -171,8 +175,7 @@ public class BaiduSearchFilter
 /// </summary>
 public class BaiduSearchRange
 {
-    [JsonPropertyName("page_time")]
-    public BaiduSearchPageTime? PageTime { get; set; }
+    [JsonPropertyName("page_time")] public BaiduSearchPageTime? PageTime { get; set; }
 }

 /// <summary>
@@ -180,11 +183,9 @@ public class BaiduSearchRange
 /// </summary>
 public class BaiduSearchPageTime
 {
-    [JsonPropertyName("gte")]
-    public string? Gte { get; set; }
+    [JsonPropertyName("gte")] public string? Gte { get; set; }

-    [JsonPropertyName("lte")]
-    public string? Lte { get; set; }
+    [JsonPropertyName("lte")] public string? Lte { get; set; }
 }

 /// <summary>
@@ -192,11 +193,9 @@ public class BaiduSearchPageTime
 /// </summary>
 public class BaiduSearchMessage
 {
-    [JsonPropertyName("role")]
-    public string Role { get; set; } = "user";
+    [JsonPropertyName("role")] public string Role { get; set; } = "user";

-    [JsonPropertyName("content")]
-    public string Content { get; set; } = "";
+    [JsonPropertyName("content")] public string Content { get; set; } = "";
 }

 /// <summary>
@@ -204,11 +203,9 @@ public class BaiduSearchMessage
 /// </summary>
 public class BaiduSearchResponse
 {
-    [JsonPropertyName("request_id")]
-    public string? RequestId { get; set; }
+    [JsonPropertyName("request_id")] public string? RequestId { get; set; }

-    [JsonPropertyName("references")]
-    public List<BaiduSearchReference>? References { get; set; }
+    [JsonPropertyName("references")] public List<BaiduSearchReference>? References { get; set; }
 }

 /// <summary>
@@ -216,23 +213,17 @@ public class BaiduSearchResponse
 /// </summary>
 public class BaiduSearchReference
 {
-    [JsonPropertyName("id")]
-    public int Id { get; set; }
+    [JsonPropertyName("id")] public int Id { get; set; }

-    [JsonPropertyName("url")]
-    public string? Url { get; set; }
+    [JsonPropertyName("url")] public string? Url { get; set; }

-    [JsonPropertyName("title")]
-    public string? Title { get; set; }
+    [JsonPropertyName("title")] public string? Title { get; set; }

-    [JsonPropertyName("date")]
-    public string? Date { get; set; }
+    [JsonPropertyName("date")] public string? Date { get; set; }

-    [JsonPropertyName("snippet")]
-    public string? Snippet { get; set; }
+    [JsonPropertyName("snippet")] public string? Snippet { get; set; }

-    [JsonPropertyName("website")]
-    public string? Website { get; set; }
+    [JsonPropertyName("website")] public string? Website { get; set; }
 }

 #endregion
--- a/Yi.Ai.Vue3/types/components.d.ts
+++ b/Yi.Ai.Vue3/types/components.d.ts
@@ -20,6 +20,7 @@ declare module 'vue' {
    ElCard: typeof import('element-plus/es')['ElCard']
    ElCollapse: typeof import('element-plus/es')['ElCollapse']
    ElCollapseItem: typeof import('element-plus/es')['ElCollapseItem']
+    ElCollapseTransition: typeof import('element-plus/es')['ElCollapseTransition']
    ElContainer: typeof import('element-plus/es')['ElContainer']
    ElDialog: typeof import('element-plus/es')['ElDialog']
    ElDivider: typeof import('element-plus/es')['ElDivider']
@@ -34,8 +35,10 @@ declare module 'vue' {
    ElMain: typeof import('element-plus/es')['ElMain']
    ElMenu: typeof import('element-plus/es')['ElMenu']
    ElMenuItem: typeof import('element-plus/es')['ElMenuItem']
+    ElOption: typeof import('element-plus/es')['ElOption']
    ElScrollbar: typeof import('element-plus/es')['ElScrollbar']
    ElSegmented: typeof import('element-plus/es')['ElSegmented']
+    ElSelect: typeof import('element-plus/es')['ElSelect']
    ElSubMenu: typeof import('element-plus/es')['ElSubMenu']
    ElTabPane: typeof import('element-plus/es')['ElTabPane']
    ElTabs: typeof import('element-plus/es')['ElTabs']