From 2b465722df0a2cec34a52c84580947c12f859430 Mon Sep 17 00:00:00 2001 From: Jicheng Lu Date: Tue, 24 Jun 2025 22:09:13 -0500 Subject: [PATCH] use stream message from state and add token usage --- .../Conversations/Dtos/ChatResponseDto.cs | 3 ++ .../Settings/ConversationSetting.cs | 1 - .../Routing/Reasoning/InstructExecutor.cs | 5 +-- .../BotSharp.Core/Routing/RoutingService.cs | 5 +-- .../Hooks/ChatHubConversationHook.cs | 9 ++++-- .../Providers/Chat/ChatCompletionProvider.cs | 31 +++++++++++++++++++ 6 files changed, 47 insertions(+), 7 deletions(-) diff --git a/src/Infrastructure/BotSharp.Abstraction/Conversations/Dtos/ChatResponseDto.cs b/src/Infrastructure/BotSharp.Abstraction/Conversations/Dtos/ChatResponseDto.cs index b391bc7b..19b5e8b2 100644 --- a/src/Infrastructure/BotSharp.Abstraction/Conversations/Dtos/ChatResponseDto.cs +++ b/src/Infrastructure/BotSharp.Abstraction/Conversations/Dtos/ChatResponseDto.cs @@ -35,6 +35,9 @@ public class ChatResponseDto : InstructResult [JsonPropertyName("has_message_files")] public bool HasMessageFiles { get; set; } + [JsonPropertyName("is_streaming")] + public bool IsStreaming { get; set; } + [JsonPropertyName("created_at")] public DateTime CreatedAt { get; set; } = DateTime.UtcNow; } diff --git a/src/Infrastructure/BotSharp.Abstraction/Conversations/Settings/ConversationSetting.cs b/src/Infrastructure/BotSharp.Abstraction/Conversations/Settings/ConversationSetting.cs index 980569dc..c4f131dc 100644 --- a/src/Infrastructure/BotSharp.Abstraction/Conversations/Settings/ConversationSetting.cs +++ b/src/Infrastructure/BotSharp.Abstraction/Conversations/Settings/ConversationSetting.cs @@ -12,7 +12,6 @@ public class ConversationSetting public bool EnableContentLog { get; set; } public bool EnableStateLog { get; set; } public bool EnableTranslationMemory { get; set; } - public bool EnableStreaming { get; set; } public CleanConversationSetting CleanSetting { get; set; } = new(); public RateLimitSetting RateLimit { get; set; } = new(); } diff --git a/src/Infrastructure/BotSharp.Core/Routing/Reasoning/InstructExecutor.cs b/src/Infrastructure/BotSharp.Core/Routing/Reasoning/InstructExecutor.cs index 42a3d652..18d87cb9 100644 --- a/src/Infrastructure/BotSharp.Core/Routing/Reasoning/InstructExecutor.cs +++ b/src/Infrastructure/BotSharp.Core/Routing/Reasoning/InstructExecutor.cs @@ -57,8 +57,9 @@ public class InstructExecutor : IExecutor } else { - var convSettings = _services.GetRequiredService(); - var ret = await routing.InvokeAgent(agentId, dialogs, convSettings.EnableStreaming); + var state = _services.GetRequiredService(); + var useStreamMsg = state.GetState("use_stream_message"); + var ret = await routing.InvokeAgent(agentId, dialogs, bool.TryParse(useStreamMsg, out var useStream) && useStream); } var response = dialogs.Last(); diff --git a/src/Infrastructure/BotSharp.Core/Routing/RoutingService.cs b/src/Infrastructure/BotSharp.Core/Routing/RoutingService.cs index bc9e1da2..e4764b92 100644 --- a/src/Infrastructure/BotSharp.Core/Routing/RoutingService.cs +++ b/src/Infrastructure/BotSharp.Core/Routing/RoutingService.cs @@ -51,8 +51,9 @@ public partial class RoutingService : IRoutingService } else { - var convSettings = _services.GetRequiredService(); - var ret = await routing.InvokeAgent(agentId, dialogs, convSettings.EnableStreaming); + var state = _services.GetRequiredService(); + var useStreamMsg = state.GetState("use_stream_message"); + var ret = await routing.InvokeAgent(agentId, dialogs, bool.TryParse(useStreamMsg, out var useStream) && useStream); } var response = dialogs.Last(); diff --git a/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/ChatHubConversationHook.cs b/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/ChatHubConversationHook.cs index 79604c98..7bc4600a 100644 --- a/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/ChatHubConversationHook.cs +++ b/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/ChatHubConversationHook.cs @@ -105,7 +105,7 @@ public class ChatHubConversationHook : ConversationHookBase public override async Task OnResponseGenerated(RoleDialogModel message) { - if (!AllowSendingMessage() || message.IsStreaming) return; + if (!AllowSendingMessage()) return; var conv = _services.GetRequiredService(); var state = _services.GetRequiredService(); @@ -118,6 +118,7 @@ public class ChatHubConversationHook : ConversationHookBase RichContent = message.SecondaryRichContent ?? message.RichContent, Data = message.Data, States = state.GetStates(), + IsStreaming = message.IsStreaming, Sender = new() { FirstName = "AI", @@ -133,7 +134,11 @@ public class ChatHubConversationHook : ConversationHookBase SenderAction = SenderActionEnum.TypingOff }; - await GenerateSenderAction(conv.ConversationId, action); + if (!message.IsStreaming) + { + await GenerateSenderAction(conv.ConversationId, action); + } + await ReceiveAssistantMessage(conv.ConversationId, json); await base.OnResponseGenerated(message); } diff --git a/src/Plugins/BotSharp.Plugin.OpenAI/Providers/Chat/ChatCompletionProvider.cs b/src/Plugins/BotSharp.Plugin.OpenAI/Providers/Chat/ChatCompletionProvider.cs index bb8c484a..16e56b25 100644 --- a/src/Plugins/BotSharp.Plugin.OpenAI/Providers/Chat/ChatCompletionProvider.cs +++ b/src/Plugins/BotSharp.Plugin.OpenAI/Providers/Chat/ChatCompletionProvider.cs @@ -1,6 +1,9 @@ +using Azure; using BotSharp.Abstraction.Hooks; using BotSharp.Core.Infrastructures.Streams; using BotSharp.Core.Observables.Queues; +using BotSharp.Plugin.OpenAI.Models.Realtime; +using Fluid; using OpenAI.Chat; namespace BotSharp.Plugin.OpenAI.Providers.Chat; @@ -190,6 +193,13 @@ public class ChatCompletionProvider : IChatCompletion var hub = _services.GetRequiredService(); var messageId = conversations.LastOrDefault()?.MessageId ?? string.Empty; + var contentHooks = _services.GetHooks(agent.Id); + // Before chat completion hook + foreach (var hook in contentHooks) + { + await hook.BeforeGenerating(agent, conversations); + } + hub.Push(new() { ServiceProvider = _services, @@ -201,8 +211,11 @@ public class ChatCompletionProvider : IChatCompletion } }); + using var textStream = new RealtimeTextStream(); var toolCalls = new List(); + ChatTokenUsage? tokenUsage = null; + var responseMessage = new RoleDialogModel(AgentRole.Assistant, string.Empty) { CurrentAgentId = agent.Id, @@ -211,6 +224,8 @@ public class ChatCompletionProvider : IChatCompletion await foreach (var choice in chatClient.CompleteChatStreamingAsync(messages, options)) { + tokenUsage = choice.Usage; + if (!choice.ToolCallUpdates.IsNullOrEmpty()) { toolCalls.AddRange(choice.ToolCallUpdates); @@ -281,6 +296,22 @@ public class ChatCompletionProvider : IChatCompletion Data = responseMessage }); + + var inputTokenDetails = tokenUsage?.InputTokenDetails; + // After chat completion hook + foreach (var hook in contentHooks) + { + await hook.AfterGenerated(responseMessage, new TokenStatsModel + { + Prompt = prompt, + Provider = Provider, + Model = _model, + TextInputTokens = (tokenUsage?.InputTokenCount ?? 0) - (inputTokenDetails?.CachedTokenCount ?? 0), + CachedTextInputTokens = inputTokenDetails?.CachedTokenCount ?? 0, + TextOutputTokens = tokenUsage?.OutputTokenCount ?? 0 + }); + } + return responseMessage; }