From be7b5a98495a6a55c75b42bc72d48823aeaeb9ba Mon Sep 17 00:00:00 2001
From: Haiping Chen <hchen@lessen.com>
Date: Thu, 14 Mar 2024 19:52:58 -0500
Subject: [PATCH] Fix InvokeFunction name.

---
 .../BotSharp.Abstraction/Routing/IRoutingService.cs         | 2 +-
 .../BotSharp.Core/Routing/Planning/NaivePlanner.cs          | 3 +++
 .../BotSharp.Core/Routing/RoutingService.InvokeFunction.cs  | 5 ++---
 .../templates/planner_prompt.naive.liquid                   | 6 +++---
 4 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/src/Infrastructure/BotSharp.Abstraction/Routing/IRoutingService.cs b/src/Infrastructure/BotSharp.Abstraction/Routing/IRoutingService.cs
index 771223ef6..74845e1f9 100644
--- a/src/Infrastructure/BotSharp.Abstraction/Routing/IRoutingService.cs
+++ b/src/Infrastructure/BotSharp.Abstraction/Routing/IRoutingService.cs
@@ -31,7 +31,7 @@ public interface IRoutingService
     List<RoutingHandlerDef> GetHandlers(Agent router);
     void ResetRecursiveCounter();
     Task<bool> InvokeAgent(string agentId, List<RoleDialogModel> dialogs);
-    Task<bool> InvokeFunction(string name, RoleDialogModel message, bool restoreOriginalFunctionName = true);
+    Task<bool> InvokeFunction(string name, RoleDialogModel message);
     Task<RoleDialogModel> InstructLoop(RoleDialogModel message);
 
     /// <summary>
diff --git a/src/Infrastructure/BotSharp.Core/Routing/Planning/NaivePlanner.cs b/src/Infrastructure/BotSharp.Core/Routing/Planning/NaivePlanner.cs
index bca5105c0..3385eff02 100644
--- a/src/Infrastructure/BotSharp.Core/Routing/Planning/NaivePlanner.cs
+++ b/src/Infrastructure/BotSharp.Core/Routing/Planning/NaivePlanner.cs
@@ -1,3 +1,4 @@
+using Amazon.Runtime.Internal.Transform;
 using BotSharp.Abstraction.Agents.Models;
 using BotSharp.Abstraction.Functions.Models;
 using BotSharp.Abstraction.Repositories.Filters;
@@ -111,9 +112,11 @@ private string GetNextStepPrompt(Agent router)
     {
         var template = router.Templates.First(x => x.Name == "planner_prompt.naive").Content;
 
+        var states = _services.GetRequiredService<IConversationStateService>();
         var render = _services.GetRequiredService<ITemplateRender>();
         return render.Render(template, new Dictionary<string, object>
         {
+            { "next_action_agent",  states.GetState("next_action_agent")}
         });
     }
 
diff --git a/src/Infrastructure/BotSharp.Core/Routing/RoutingService.InvokeFunction.cs b/src/Infrastructure/BotSharp.Core/Routing/RoutingService.InvokeFunction.cs
index 6f70e0638..77659d69e 100644
--- a/src/Infrastructure/BotSharp.Core/Routing/RoutingService.InvokeFunction.cs
+++ b/src/Infrastructure/BotSharp.Core/Routing/RoutingService.InvokeFunction.cs
@@ -4,7 +4,7 @@ namespace BotSharp.Core.Routing;
 
 public partial class RoutingService
 {
-    public async Task<bool> InvokeFunction(string name, RoleDialogModel message, bool restoreOriginalFunctionName = true)
+    public async Task<bool> InvokeFunction(string name, RoleDialogModel message)
     {
         var function = _services.GetServices<IFunctionCallback>().FirstOrDefault(x => x.Name == name);
         if (function == null)
@@ -57,8 +57,7 @@ public async Task<bool> InvokeFunction(string name, RoleDialogModel message, boo
 
         // restore original function name
         if (!message.StopCompletion && 
-            message.FunctionName != originalFunctionName &&
-            restoreOriginalFunctionName)
+            message.FunctionName != originalFunctionName)
         {
             message.FunctionName = originalFunctionName;
         }
diff --git a/src/Infrastructure/BotSharp.Core/data/agents/01fcc3e5-9af7-49e6-ad7a-a760bd12dc4a/templates/planner_prompt.naive.liquid b/src/Infrastructure/BotSharp.Core/data/agents/01fcc3e5-9af7-49e6-ad7a-a760bd12dc4a/templates/planner_prompt.naive.liquid
index 449bd4d7b..9407226cf 100644
--- a/src/Infrastructure/BotSharp.Core/data/agents/01fcc3e5-9af7-49e6-ad7a-a760bd12dc4a/templates/planner_prompt.naive.liquid
+++ b/src/Infrastructure/BotSharp.Core/data/agents/01fcc3e5-9af7-49e6-ad7a-a760bd12dc4a/templates/planner_prompt.naive.liquid
@@ -1,3 +1,3 @@
-What is the next step based on the CONVERSATION? 
-Route to the Agent that last handled the conversation if necessary. 
-If user wants to speak to customer service, use function human_intervention_needed.
\ No newline at end of file
+What is the next step based on the CONVERSATION?                                   
+If user is replying with agent's question, you must route back to the previous agent {{ next_action_agent }}. 
+If user wants to speak to customer service, use function human_intervention_needed. 
\ No newline at end of file