From 9e2d2807a126535f37783639f4aa9c7d4124bf5a Mon Sep 17 00:00:00 2001
From: Joanna Ren <101223@smsassist.com>
Date: Wed, 4 Sep 2024 10:16:05 -0500
Subject: [PATCH] add DDL to planner

---
 .../Functions/PrimaryStagePlanFn.cs           |  5 +++-
 .../Functions/SummaryPlanFn.cs                | 24 ++++++++++++++++++-
 .../instructions/instruction.liquid           |  2 +-
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/Plugins/BotSharp.Plugin.Planner/Functions/PrimaryStagePlanFn.cs b/src/Plugins/BotSharp.Plugin.Planner/Functions/PrimaryStagePlanFn.cs
index a12cafed3..585dc1746 100644
--- a/src/Plugins/BotSharp.Plugin.Planner/Functions/PrimaryStagePlanFn.cs
+++ b/src/Plugins/BotSharp.Plugin.Planner/Functions/PrimaryStagePlanFn.cs
@@ -39,7 +39,10 @@ public async Task<bool> Execute(RoleDialogModel message)
         };
         var retrievalMessage = new RoleDialogModel(AgentRole.User, task.Requirements)
         {
-            FunctionArgs = JsonSerializer.Serialize(msg),
+            FunctionArgs = JsonSerializer.Serialize(new ExtractedKnowledge 
+            { 
+                Question = task.Requirements
+            }),
             KnowledgeConfidence = 0.1f,
             Content = ""
         };
diff --git a/src/Plugins/BotSharp.Plugin.Planner/Functions/SummaryPlanFn.cs b/src/Plugins/BotSharp.Plugin.Planner/Functions/SummaryPlanFn.cs
index 098edd533..4ec2a45d4 100644
--- a/src/Plugins/BotSharp.Plugin.Planner/Functions/SummaryPlanFn.cs
+++ b/src/Plugins/BotSharp.Plugin.Planner/Functions/SummaryPlanFn.cs
@@ -5,6 +5,8 @@
 using BotSharp.Core.Infrastructures;
 using BotSharp.Plugin.Planner.TwoStaging.Models;
 using Microsoft.Extensions.Logging;
+using BotSharp.Abstraction.Routing;
+using BotSharp.Core.Agents.Services;
 
 namespace BotSharp.Plugin.Planner.Functions;
 
@@ -23,12 +25,31 @@ public SummaryPlanFn(IServiceProvider services, ILogger<PrimaryStagePlanFn> logg
 
     public async Task<bool> Execute(RoleDialogModel message)
     {
+        var fn = _services.GetRequiredService<IRoutingService>();
+        var agentService = _services.GetRequiredService<IAgentService>();
+        var currentAgent = await agentService.LoadAgent(message.CurrentAgentId);
         //debug
         var state = _services.GetRequiredService<IConversationStateService>();
         state.SetState("max_tokens", "4096");
 
         var task = state.GetState("requirement_detail");
 
+        //get DDL
+        var steps = message.Content.JsonArrayContent<SecondStagePlan>();
+
+        //get all the related tables
+        List<string> allTables = new List<string>();
+        foreach (var step in steps)
+        {
+            allTables.AddRange(step.Tables);
+        }
+        message.Data = allTables.Distinct().ToList();
+
+        //get table DDL and stores in content
+        var msg2 = RoleDialogModel.From(message);
+        await fn.InvokeFunction("get_table_definition", msg2);
+        message.SecondaryContent = msg2.Content;
+
         // summarize and generate query
         var summaryPlanningPrompt = await GetPlanSummaryPrompt(task, message);
         _logger.LogInformation(summaryPlanningPrompt);
@@ -37,7 +58,8 @@ public async Task<bool> Execute(RoleDialogModel message)
             Id = BuiltInAgentId.Planner,
             Name = "planner_summary",
             Instruction = summaryPlanningPrompt,
-            TemplateDict = new Dictionary<string, object>()
+            TemplateDict = new Dictionary<string, object>(),
+            LlmConfig = currentAgent.LlmConfig
         };
         var response_summary = await GetAIResponse(plannerAgent);
 
diff --git a/src/Plugins/BotSharp.Plugin.Planner/data/agents/282a7128-69a1-44b0-878c-a9159b88f3b9/instructions/instruction.liquid b/src/Plugins/BotSharp.Plugin.Planner/data/agents/282a7128-69a1-44b0-878c-a9159b88f3b9/instructions/instruction.liquid
index f40fa6c8d..7209f7dd1 100644
--- a/src/Plugins/BotSharp.Plugin.Planner/data/agents/282a7128-69a1-44b0-878c-a9159b88f3b9/instructions/instruction.liquid
+++ b/src/Plugins/BotSharp.Plugin.Planner/data/agents/282a7128-69a1-44b0-878c-a9159b88f3b9/instructions/instruction.liquid
@@ -1,3 +1,3 @@
 Use the TwoStagePlanner approach to plan the overall implementation steps, call plan_primary_stage. 
 If need_additional_information is true, call plan_secondary_stage for the specific primary stage.
-Call plan_summary to summarize the final planning steps.
\ No newline at end of file
+You must Call plan_summary as the last step to summarize the final planning steps.
\ No newline at end of file