From 829e4403cece675ed6e6386be06fc2a7d48d5c8f Mon Sep 17 00:00:00 2001
From: Haiping Chen <101423@smsassist.com>
Date: Tue, 1 Oct 2024 16:43:20 -0500
Subject: [PATCH] sql output formatting.

---
 .../Agents/Enums/BuiltInAgentId.cs            |  2 ++
 .../BotSharp.Plugin.SqlDriver.csproj          |  7 +++---
 .../Functions/ExecuteQueryFn.cs               | 25 +++++++++++++++++++
 .../Models/ExecuteQueryArgs.cs                |  5 ++++
 .../templates/query_result_formatting.liquid  |  1 +
 5 files changed, 37 insertions(+), 3 deletions(-)
 create mode 100644 src/Plugins/BotSharp.Plugin.SqlDriver/data/agents/beda4c12-e1ec-4b4b-b328-3df4a6687c4f/templates/query_result_formatting.liquid

diff --git a/src/Infrastructure/BotSharp.Abstraction/Agents/Enums/BuiltInAgentId.cs b/src/Infrastructure/BotSharp.Abstraction/Agents/Enums/BuiltInAgentId.cs
index 2c95818f..55436ae2 100644
--- a/src/Infrastructure/BotSharp.Abstraction/Agents/Enums/BuiltInAgentId.cs
+++ b/src/Infrastructure/BotSharp.Abstraction/Agents/Enums/BuiltInAgentId.cs
@@ -36,4 +36,6 @@ public class BuiltInAgentId
     /// Plan feasible implementation steps for complex problems
     /// </summary>
     public const string Planner = "282a7128-69a1-44b0-878c-a9159b88f3b9";
+
+    public const string SqlDriver = "beda4c12-e1ec-4b4b-b328-3df4a6687c4f";
 }
diff --git a/src/Plugins/BotSharp.Plugin.SqlDriver/BotSharp.Plugin.SqlDriver.csproj b/src/Plugins/BotSharp.Plugin.SqlDriver/BotSharp.Plugin.SqlDriver.csproj
index 02bca766..4f93597a 100644
--- a/src/Plugins/BotSharp.Plugin.SqlDriver/BotSharp.Plugin.SqlDriver.csproj
+++ b/src/Plugins/BotSharp.Plugin.SqlDriver/BotSharp.Plugin.SqlDriver.csproj
@@ -11,11 +11,8 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <Compile Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\**" />
     <Compile Remove="packages\**" />
-    <EmbeddedResource Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\**" />
     <EmbeddedResource Remove="packages\**" />
-    <None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\**" />
     <None Remove="packages\**" />
   </ItemGroup>
 
@@ -32,6 +29,7 @@
     <None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\functions\sql_insert.json" />
     <None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\functions\sql_select.json" />
     <None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\instructions\instruction.liquid" />
+    <None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\query_result_formatting.liquid" />
   </ItemGroup>
 
   <ItemGroup>
@@ -71,6 +69,9 @@
     <Content Include="data\agents\6745151e-6d46-4a02-8de4-1c4f21c7da95\templates\sql_executor.fn.liquid">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
     </Content>
+    <Content Include="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\query_result_formatting.liquid">
+      <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
+    </Content>
   </ItemGroup>
 
   <ItemGroup>
diff --git a/src/Plugins/BotSharp.Plugin.SqlDriver/Functions/ExecuteQueryFn.cs b/src/Plugins/BotSharp.Plugin.SqlDriver/Functions/ExecuteQueryFn.cs
index d8ae5097..165b2211 100644
--- a/src/Plugins/BotSharp.Plugin.SqlDriver/Functions/ExecuteQueryFn.cs
+++ b/src/Plugins/BotSharp.Plugin.SqlDriver/Functions/ExecuteQueryFn.cs
@@ -1,3 +1,5 @@
+using BotSharp.Abstraction.Agents.Enums;
+using BotSharp.Core.Infrastructures;
 using BotSharp.Plugin.SqlDriver.Models;
 using Dapper;
 using Microsoft.Data.SqlClient;
@@ -30,6 +32,29 @@ public class ExecuteQueryFn : IFunctionCallback
         };
 
         message.Content = JsonSerializer.Serialize(results);
+
+        if (args.FormattingResult)
+        {
+            var conv = _services.GetRequiredService<IConversationService>();
+            var sqlAgent = await _services.GetRequiredService<IAgentService>().LoadAgent(BuiltInAgentId.SqlDriver);
+            var prompt = sqlAgent.Templates.FirstOrDefault(x => x.Name == "query_result_formatting");
+
+            var completion = CompletionProvider.GetChatCompletion(_services,
+                provider: sqlAgent.LlmConfig.Provider,
+                model: sqlAgent.LlmConfig.Model);
+
+            var result = await completion.GetChatCompletions(new Agent
+            {
+                Id = sqlAgent.Id,
+                Instruction = prompt.Content,
+            }, new List<RoleDialogModel>
+            {
+                new RoleDialogModel(AgentRole.User, message.Content)
+            });
+
+            message.Content = result.Content;
+        }
+
         return true;
     }
 
diff --git a/src/Plugins/BotSharp.Plugin.SqlDriver/Models/ExecuteQueryArgs.cs b/src/Plugins/BotSharp.Plugin.SqlDriver/Models/ExecuteQueryArgs.cs
index b1531e64..ac2279ea 100644
--- a/src/Plugins/BotSharp.Plugin.SqlDriver/Models/ExecuteQueryArgs.cs
+++ b/src/Plugins/BotSharp.Plugin.SqlDriver/Models/ExecuteQueryArgs.cs
@@ -6,4 +6,9 @@ public class ExecuteQueryArgs
 {
     [JsonPropertyName("sql_statements")]
     public string[] SqlStatements { get; set; } = [];
+
+    /// <summary>
+    /// Beautifying query result
+    /// </summary>
+    public bool FormattingResult { get; set; }
 }
diff --git a/src/Plugins/BotSharp.Plugin.SqlDriver/data/agents/beda4c12-e1ec-4b4b-b328-3df4a6687c4f/templates/query_result_formatting.liquid b/src/Plugins/BotSharp.Plugin.SqlDriver/data/agents/beda4c12-e1ec-4b4b-b328-3df4a6687c4f/templates/query_result_formatting.liquid
new file mode 100644
index 00000000..0ef43495
--- /dev/null
+++ b/src/Plugins/BotSharp.Plugin.SqlDriver/data/agents/beda4c12-e1ec-4b4b-b328-3df4a6687c4f/templates/query_result_formatting.liquid
@@ -0,0 +1 @@
+Output in human readable format.
\ No newline at end of file