From a768dc370109b00834bdee48ab639dbcb2995109 Mon Sep 17 00:00:00 2001
From: Jicheng Lu <103353@smsassist.com>
Date: Fri, 26 Jan 2024 15:47:06 -0600
Subject: [PATCH] add conv state log

---
 .../Models/ConversationStateLogModel.cs            | 11 +++++++++++
 .../Loggers/Models/StreamingLogModel.cs            |  2 ++
 .../Hooks/ChatHubConversationHook.cs               | 14 ++++++++++++++
 .../Hooks/StreamingLogHook.cs                      |  9 +++++----
 4 files changed, 32 insertions(+), 4 deletions(-)
 create mode 100644 src/Infrastructure/BotSharp.Abstraction/Conversations/Models/ConversationStateLogModel.cs

diff --git a/src/Infrastructure/BotSharp.Abstraction/Conversations/Models/ConversationStateLogModel.cs b/src/Infrastructure/BotSharp.Abstraction/Conversations/Models/ConversationStateLogModel.cs
new file mode 100644
index 00000000..de80b485
--- /dev/null
+++ b/src/Infrastructure/BotSharp.Abstraction/Conversations/Models/ConversationStateLogModel.cs
@@ -0,0 +1,11 @@
+namespace BotSharp.Abstraction.Conversations.Models;
+
+public class ConversationStateLogModel
+{
+    [JsonPropertyName("conversation_id")]
+    public string ConvsersationId { get; set; }
+    [JsonPropertyName("states")]
+    public string States { get; set; }
+    [JsonPropertyName("created_at")]
+    public DateTime CreateTime { get; set; }
+}
diff --git a/src/Infrastructure/BotSharp.Abstraction/Loggers/Models/StreamingLogModel.cs b/src/Infrastructure/BotSharp.Abstraction/Loggers/Models/StreamingLogModel.cs
index d361bf59..51941138 100644
--- a/src/Infrastructure/BotSharp.Abstraction/Loggers/Models/StreamingLogModel.cs
+++ b/src/Infrastructure/BotSharp.Abstraction/Loggers/Models/StreamingLogModel.cs
@@ -4,6 +4,8 @@ public class StreamingLogModel
 {
     [JsonPropertyName("conversation_id")]
     public string ConversationId { get; set; }
+    [JsonPropertyName("name")]
+    public string? Name { get; set; }
 
     [JsonPropertyName("content")]
     public string Content { get; set; }
diff --git a/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/ChatHubConversationHook.cs b/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/ChatHubConversationHook.cs
index 1b9a0089..bad8cf2b 100644
--- a/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/ChatHubConversationHook.cs
+++ b/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/ChatHubConversationHook.cs
@@ -100,6 +100,7 @@ public class ChatHubConversationHook : ConversationHookBase
     public override async Task OnResponseGenerated(RoleDialogModel message)
     {
         var conv = _services.GetRequiredService<IConversationService>();
+        var state = _services.GetRequiredService<IConversationStateService>();
 
         var json = JsonSerializer.Serialize(new ChatResponseModel()
         {
@@ -115,7 +116,20 @@ public class ChatHubConversationHook : ConversationHookBase
             }
         }, _serializerOptions);
         await _chatHub.Clients.User(_user.Id).SendAsync("OnMessageReceivedFromAssistant", json);
+        await _chatHub.Clients.User(_user.Id).SendAsync("OnConversateStatesGenerated", BuildConversationStates(conv.ConversationId, state.GetStates()));
 
         await base.OnResponseGenerated(message);
     }
+
+    private string BuildConversationStates(string conversationId, Dictionary<string, string> states)
+    {
+        var model = new ConversationStateLogModel
+        {
+            ConvsersationId = conversationId,
+            States = JsonSerializer.Serialize(states, _serializerOptions),
+            CreateTime = DateTime.UtcNow
+        };
+
+        return JsonSerializer.Serialize(model, _serializerOptions);
+    }
 }
diff --git a/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/StreamingLogHook.cs b/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/StreamingLogHook.cs
index bdf43c13..8b0aa2cd 100644
--- a/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/StreamingLogHook.cs
+++ b/src/Plugins/BotSharp.Plugin.ChatHub/Hooks/StreamingLogHook.cs
@@ -37,7 +37,7 @@ public class StreamingLogHook : ConversationHookBase, IContentGeneratingHook
     {
         var conversationId = _state.GetConversationId();
         var log = $"MessageId: {message.MessageId} ==>\r\n{message.Role}: {message.Content}";
-        await _chatHub.Clients.User(_user.Id).SendAsync("OnContentLogGenerated", BuildLog(conversationId, log));
+        await _chatHub.Clients.User(_user.Id).SendAsync("OnContentLogGenerated", BuildLog(conversationId, _user.UserName, log));
     }
 
     public async Task BeforeGenerating(Agent agent, List<RoleDialogModel> conversations)
@@ -64,20 +64,21 @@ public class StreamingLogHook : ConversationHookBase, IContentGeneratingHook
         var conversationId = _state.GetConversationId();
         var agent = await agentService.LoadAgent(message.CurrentAgentId);
 
-        await _chatHub.Clients.User(_user.Id).SendAsync("OnContentLogGenerated", BuildLog(conversationId, tokenStats.Prompt));
+        await _chatHub.Clients.User(_user.Id).SendAsync("OnContentLogGenerated", BuildLog(conversationId, agent?.Name, tokenStats.Prompt));
 
         var log = message.Role == AgentRole.Function ?
                 $"[{agent?.Name}]: {message.FunctionName}({message.FunctionArgs})" :
                 $"[{agent?.Name}]: {message.Content}";
         log += $"\r\n<== MessageId: {message.MessageId}";
-        await _chatHub.Clients.User(_user.Id).SendAsync("OnContentLogGenerated", BuildLog(conversationId, log));
+        await _chatHub.Clients.User(_user.Id).SendAsync("OnContentLogGenerated", BuildLog(conversationId, agent?.Name, log));
     }
 
-    private string BuildLog(string conversationId, string content)
+    private string BuildLog(string conversationId, string? name, string content)
     {
         var log = new StreamingLogModel
         {
             ConversationId = conversationId,
+            Name = name,
             Content = content,
             CreateTime = DateTime.UtcNow
         };