use stream message from state and add token usage

This commit is contained in:
Jicheng Lu 2025-06-24 22:09:13 -05:00
parent d754e5671a
commit 2b465722df
6 changed files with 47 additions and 7 deletions

View file

@ -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;
}

View file

@ -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();
}

View file

@ -57,8 +57,9 @@ public class InstructExecutor : IExecutor
}
else
{
var convSettings = _services.GetRequiredService<ConversationSetting>();
var ret = await routing.InvokeAgent(agentId, dialogs, convSettings.EnableStreaming);
var state = _services.GetRequiredService<IConversationStateService>();
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();

View file

@ -51,8 +51,9 @@ public partial class RoutingService : IRoutingService
}
else
{
var convSettings = _services.GetRequiredService<ConversationSetting>();
var ret = await routing.InvokeAgent(agentId, dialogs, convSettings.EnableStreaming);
var state = _services.GetRequiredService<IConversationStateService>();
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();

View file

@ -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<IConversationService>();
var state = _services.GetRequiredService<IConversationStateService>();
@ -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);
}

View file

@ -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<MessageHub>();
var messageId = conversations.LastOrDefault()?.MessageId ?? string.Empty;
var contentHooks = _services.GetHooks<IContentGeneratingHook>(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<StreamingChatToolCallUpdate>();
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;
}