From 61c1f9623b43ce8617289f91b758ee6bafec30fd Mon Sep 17 00:00:00 2001
From: Haiping Chen <haiping008@gmail.com>
Date: Tue, 12 Dec 2023 08:16:34 -0600
Subject: [PATCH] Fix StopCompletion issue for hook.OnMessageReceived.

---
 .../Conversations/IConversationService.cs     |  2 +-
 .../ConversationService.SendMessage.cs        | 37 +++++++++++--------
 .../Routing/RoutingService.InvokeFunction.cs  |  1 +
 3 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/src/Infrastructure/BotSharp.Abstraction/Conversations/IConversationService.cs b/src/Infrastructure/BotSharp.Abstraction/Conversations/IConversationService.cs
index cc19a3d4..d41a8831 100644
--- a/src/Infrastructure/BotSharp.Abstraction/Conversations/IConversationService.cs
+++ b/src/Infrastructure/BotSharp.Abstraction/Conversations/IConversationService.cs
@@ -29,7 +29,7 @@ public interface IConversationService
         Func<RoleDialogModel, Task> onFunctionExecuting,
         Func<RoleDialogModel, Task> onFunctionExecuted);
 
-    List<RoleDialogModel> GetDialogHistory(int lastCount = 20);
+    List<RoleDialogModel> GetDialogHistory(int lastCount = 50);
     Task CleanHistory(string agentId);
 
     Task CallFunctions(RoleDialogModel msg);
diff --git a/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationService.SendMessage.cs b/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationService.SendMessage.cs
index adf84334..19ccad1b 100644
--- a/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationService.SendMessage.cs
+++ b/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationService.SendMessage.cs
@@ -32,8 +32,12 @@ public partial class ConversationService
 
         _storage.Append(_conversationId, message);
 
+        var statistics = _services.GetRequiredService<ITokenStatistics>();
         var hooks = _services.GetServices<IConversationHook>().ToList();
 
+        RoleDialogModel response = message;
+        bool stopCompletion = false;
+
         // Before chat completion hook
         foreach (var hook in hooks)
         {
@@ -45,27 +49,26 @@ public partial class ConversationService
             // Interrupted by hook
             if (message.StopCompletion)
             {
-                await onMessageReceived(message);
-                _storage.Append(_conversationId, message);
-                return true;
+                stopCompletion = true;
             }
         }
 
-        // Routing with reasoning
-        var routing = _services.GetRequiredService<IRoutingService>();
-        var settings = _services.GetRequiredService<RoutingSettings>();
+        if (!stopCompletion)
+        {
+            // Routing with reasoning
+            var routing = _services.GetRequiredService<IRoutingService>();
+            var settings = _services.GetRequiredService<RoutingSettings>();
 
-        var response = agentId == settings.RouterId ?
-            await routing.InstructLoop(message) :
-            await routing.ExecuteDirectly(agent, message);
+            response = agentId == settings.RouterId ?
+                await routing.InstructLoop(message) :
+                await routing.ExecuteDirectly(agent, message);
+
+            routing.ResetRecursiveCounter();
+        }
 
         await HandleAssistantMessage(response, onMessageReceived);
-
-        var statistics = _services.GetRequiredService<ITokenStatistics>();
         statistics.PrintStatistics();
 
-        routing.ResetRecursiveCounter();
-
         return true;
     }
 
@@ -131,7 +134,11 @@ public partial class ConversationService
 
         // Add to dialog history
         _storage.Append(_conversationId, response);
-        var conversation = _services.GetRequiredService<IConversationService>();
-        var updatedConversation = await conversation.UpdateConversationTitle(_conversationId, response.Instruction.Reason);
+
+        if (response.Instruction != null)
+        {
+            var conversation = _services.GetRequiredService<IConversationService>();
+            var updatedConversation = await conversation.UpdateConversationTitle(_conversationId, response.Instruction.Reason);
+        }
     }
 }
diff --git a/src/Infrastructure/BotSharp.Core/Routing/RoutingService.InvokeFunction.cs b/src/Infrastructure/BotSharp.Core/Routing/RoutingService.InvokeFunction.cs
index 81e8ad90..ba937d2d 100644
--- a/src/Infrastructure/BotSharp.Core/Routing/RoutingService.InvokeFunction.cs
+++ b/src/Infrastructure/BotSharp.Core/Routing/RoutingService.InvokeFunction.cs
@@ -10,6 +10,7 @@ public partial class RoutingService
         if (function == null) return false;
 
         message.FunctionName = name;
+        message.Role = AgentRole.Function;
         return await function.Execute(message);
     }
 }