Skip to content

Commit 713c58b

Browse files
Oceania2018Oceania2018
authored
Merge pull request #800 from Joannall/master
add sql planner
2 parents cd7fdf4 + 80e292a commit 713c58b

File tree

27 files changed

+690
-23
lines changed

27 files changed

+690
-23
lines changed

‎src/Infrastructure/BotSharp.Core/data/agents/01fcc3e5-9af7-49e6-ad7a-a760bd12dc4a/agent.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
},
1717
{
1818
"type": "planner",
19-
"field": "Two-Stage-Planner"
19+
"field": "SQL-Planner"
2020
}
2121
]
2222
}

‎src/Plugins/BotSharp.Plugin.Planner/BotSharp.Plugin.Planner.csproj‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,36 @@
6969
<Content Include="data\agents\6745151e-6d46-4a02-8de4-1c4f21c7da95\templates\util-planner-plan_summary.fn.liquid">
7070
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
7171
</Content>
72+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\agent.json">
73+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
74+
</Content>
75+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\functions\plan_primary_stage.json">
76+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
77+
</Content>
78+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\functions\plan_secondary_stage.json">
79+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
80+
</Content>
81+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\functions\sql_review.json">
82+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
83+
</Content>
84+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\functions\sql_generation.json">
85+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
86+
</Content>
87+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\instructions\instruction.liquid">
88+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
89+
</Content>
90+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\templates\two_stage.1st.plan.liquid">
91+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
92+
</Content>
93+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\templates\two_stage.2nd.plan.liquid">
94+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
95+
</Content>
96+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\templates\two_stage.next.liquid">
97+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
98+
</Content>
99+
<Content Include="data\agents\da7aad2c-8112-48a2-ab7b-1f87da524741\templates\two_stage.summarize.liquid">
100+
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
101+
</Content>
72102
</ItemGroup>
73103

74104
<ItemGroup>

‎src/Plugins/BotSharp.Plugin.Planner/Enums/PlannerAgentId.cs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,8 @@
1-
using Microsoft.AspNetCore.Http;
2-
31
namespace BotSharp.Plugin.Planner.Enums;
42

53
public class PlannerAgentId
64
{
75
public const string TwoStagePlanner = "282a7128-69a1-44b0-878c-a9159b88f3b9";
86
public const string SequentialPlanner = "3e75e818-a139-48a8-9e22-4662548c13a3";
7+
public const string SqlPlanner = "da7aad2c-8112-48a2-ab7b-1f87da524741";
98
}

‎src/Plugins/BotSharp.Plugin.Planner/Functions/SummaryPlanFn.cs‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -68,15 +68,6 @@ public async Task<bool> Execute(RoleDialogModel message)
6868
var summary = await GetAiResponse(plannerAgent);
6969
message.Content = summary.Content;
7070

71-
// Emit event if the sql statement is generated by planner
72-
var args = JsonSerializer.Deserialize<SummaryPlan>(message.FunctionArgs);
73-
if (args != null && !args.IsSqlTemplate && args.ContainsSqlStatements)
74-
{
75-
await HookEmitter.Emit<IPlanningHook>(_services, async hook =>
76-
await hook.OnSourceCodeGenerated(nameof(TwoStageTaskPlanner), message, "sql")
77-
);
78-
}
79-
8071
await HookEmitter.Emit<IPlanningHook>(_services, async hook =>
8172
await hook.OnPlanningCompleted(nameof(TwoStageTaskPlanner), message)
8273
);

‎src/Plugins/BotSharp.Plugin.Planner/PlannerPlugin.cs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using BotSharp.Plugin.Planner.Sequential;
2+
using BotSharp.Plugin.Planner.SqlGeneration;
23
using BotSharp.Plugin.Planner.TwoStaging;
34

45
namespace BotSharp.Plugin.Planner;
@@ -16,13 +17,15 @@ public class PlannerPlugin : IBotSharpPlugin
1617
public string[] AgentIds =>
1718
[
1819
PlannerAgentId.TwoStagePlanner,
19-
PlannerAgentId.SequentialPlanner
20+
PlannerAgentId.SequentialPlanner,
21+
PlannerAgentId.SqlPlanner
2022
];
2123

2224
public void RegisterDI(IServiceCollection services, IConfiguration config)
2325
{
2426
services.AddScoped<ITaskPlanner, SequentialPlanner>();
2527
services.AddScoped<ITaskPlanner, TwoStageTaskPlanner>();
28+
services.AddScoped<ITaskPlanner, SqlGenerationPlanner>();
2629
services.AddScoped<IAgentHook, PlannerAgentHook>();
2730
services.AddScoped<IAgentUtilityHook, PlannerUtilityHook>();
2831
}

‎src/Plugins/BotSharp.Plugin.Planner/SqlGeneration/Functions/SqlGenerationFn.cs‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
using BotSharp.Plugin.Planner.TwoStaging;
2+
using BotSharp.Plugin.Planner.TwoStaging.Models;
3+
4+
namespace BotSharp.Plugin.Planner.Functions;
5+
6+
public class SqlGenerationFn : IFunctionCallback
7+
{
8+
public string Name => "sql_generation";
9+
public string Indication => "Organizing and summarizing the final SQL statements.";
10+
11+
private readonly IServiceProvider _services;
12+
private readonly ILogger<SqlGenerationFn> _logger;
13+
14+
public SqlGenerationFn(
15+
IServiceProvider services,
16+
ILogger<SqlGenerationFn> logger)
17+
{
18+
_services = services;
19+
_logger = logger;
20+
}
21+
22+
public async Task<bool> Execute(RoleDialogModel message)
23+
{
24+
var fn = _services.GetRequiredService<IRoutingService>();
25+
var agentService = _services.GetRequiredService<IAgentService>();
26+
var states = _services.GetRequiredService<IConversationStateService>();
27+
28+
states.SetState("max_tokens", "4096");
29+
var currentAgent = await agentService.LoadAgent(message.CurrentAgentId);
30+
var taskRequirement = states.GetState("requirement_detail");
31+
32+
// Get table names
33+
var steps = states.GetState("planning_result").JsonArrayContent<SecondStagePlan>();
34+
var allTables = new List<string>();
35+
var ddlStatements = string.Empty;
36+
var domainKnowledge = states.GetState("planning_result");
37+
domainKnowledge += "\r\n" + states.GetState("domain_knowledges");
38+
var dictionaryItems = states.GetState("dictionary_items");
39+
var excelImportResult = states.GetState("excel_import_result");
40+
41+
foreach (var step in steps)
42+
{
43+
allTables.AddRange(step.Tables);
44+
}
45+
var distinctTables = allTables.Distinct().ToList();
46+
47+
var msgCopy = RoleDialogModel.From(message);
48+
msgCopy.FunctionArgs = JsonSerializer.Serialize(new
49+
{
50+
tables = distinctTables,
51+
});
52+
await fn.InvokeFunction("sql_table_definition", msgCopy);
53+
ddlStatements += "\r\n" + msgCopy.Content;
54+
states.SetState("table_ddls", ddlStatements);
55+
56+
// Summarize and generate query
57+
var prompt = await GetSummaryPlanPrompt(msgCopy, taskRequirement, domainKnowledge, dictionaryItems, ddlStatements, excelImportResult);
58+
_logger.LogInformation($"Summary plan prompt:\r\n{prompt}");
59+
60+
var plannerAgent = new Agent
61+
{
62+
Id = PlannerAgentId.TwoStagePlanner,
63+
Name = Name,
64+
Instruction = prompt,
65+
LlmConfig = currentAgent.LlmConfig
66+
};
67+
68+
var summary = await GetAiResponse(plannerAgent);
69+
message.Content = summary.Content;
70+
71+
/*await HookEmitter.Emit<IPlanningHook>(_services, async hook =>
72+
await hook.OnPlanningCompleted(nameof(TwoStageTaskPlanner), message)
73+
);*/
74+
75+
return true;
76+
}
77+
78+
private async Task<string> GetSummaryPlanPrompt(RoleDialogModel message, string taskDescription, string domainKnowledge, string dictionaryItems, string ddlStatement, string excelImportResult)
79+
{
80+
var agentService = _services.GetRequiredService<IAgentService>();
81+
var render = _services.GetRequiredService<ITemplateRender>();
82+
var knowledgeHooks = _services.GetServices<IKnowledgeHook>();
83+
84+
var agent = await agentService.GetAgent(PlannerAgentId.TwoStagePlanner);
85+
var template = agent.Templates.FirstOrDefault(x => x.Name == "two_stage.summarize")?.Content ?? string.Empty;
86+
87+
var additionalRequirements = new List<string>();
88+
await HookEmitter.Emit<IPlanningHook>(_services, async x =>
89+
{
90+
var requirement = await x.GetSummaryAdditionalRequirements(nameof(TwoStageTaskPlanner), message);
91+
additionalRequirements.Add(requirement);
92+
});
93+
94+
var globalKnowledges = new List<string>();
95+
foreach (var hook in knowledgeHooks)
96+
{
97+
var k = await hook.GetGlobalKnowledges(message);
98+
globalKnowledges.AddRange(k);
99+
}
100+
101+
return render.Render(template, new Dictionary<string, object>
102+
{
103+
{ "task_description", taskDescription },
104+
{ "summary_requirements", string.Join("\r\n", additionalRequirements) },
105+
{ "global_knowledges", globalKnowledges },
106+
{ "domain_knowledges", domainKnowledge },
107+
{ "dictionary_items", dictionaryItems },
108+
{ "table_structure", ddlStatement },
109+
{ "excel_import_result", excelImportResult }
110+
});
111+
}
112+
private async Task<RoleDialogModel> GetAiResponse(Agent plannerAgent)
113+
{
114+
var conv = _services.GetRequiredService<IConversationService>();
115+
var wholeDialogs = conv.GetDialogHistory();
116+
117+
// Append text
118+
wholeDialogs.Last().Content += "\n\nIf the table structure didn't mention auto incremental, the data field id needs to insert id manually and you need to use max(id).\nFor example, you should use SET @id = select max(id) from table;";
119+
wholeDialogs.Last().Content += "\n\nTry if you can generate a single query to fulfill the needs.";
120+
121+
var completion = CompletionProvider.GetChatCompletion(_services,
122+
provider: plannerAgent.LlmConfig.Provider,
123+
model: plannerAgent.LlmConfig.Model);
124+
125+
return await completion.GetChatCompletions(plannerAgent, wholeDialogs);
126+
}
127+
}

‎src/Plugins/BotSharp.Plugin.Planner/SqlGeneration/Functions/SqlReviewFn.cs‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
using BotSharp.Plugin.Planner.SqlGeneration.Models;
2+
using BotSharp.Plugin.Planner.TwoStaging;
3+
using BotSharp.Plugin.Planner.TwoStaging.Models;
4+
5+
namespace BotSharp.Plugin.Planner.SqlGeneration.Functions;
6+
7+
public class SqlReviewFn : IFunctionCallback
8+
{
9+
public string Name => "sql_review";
10+
public string Indication => "Currently reviewing SQL statement";
11+
12+
private readonly IServiceProvider _services;
13+
private readonly ILogger<SqlReviewFn> _logger;
14+
15+
public SqlReviewFn(
16+
IServiceProvider services,
17+
ILogger<SqlReviewFn> logger)
18+
{
19+
_services = services;
20+
_logger = logger;
21+
}
22+
23+
public async Task<bool> Execute(RoleDialogModel message)
24+
{
25+
var args = JsonSerializer.Deserialize<SqlReviewArgs>(message.FunctionArgs);
26+
if (!message.Content.StartsWith("```sql"))
27+
{
28+
message.Content = $"```sql\r\n{args.SqlStatement}\r\n```";
29+
}
30+
if (args != null && !args.IsSqlTemplate && args.ContainsSqlStatements)
31+
{
32+
await HookEmitter.Emit<IPlanningHook>(_services, async hook =>
33+
await hook.OnSourceCodeGenerated(nameof(TwoStageTaskPlanner), message, "sql")
34+
);
35+
}
36+
return true;
37+
}
38+
}

‎src/Plugins/BotSharp.Plugin.Planner/SqlGeneration/Models/FirstStagePlan.cs‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
namespace BotSharp.Plugin.Planner.SqlGeneration.Models;
2+
3+
public class FirstStagePlan
4+
{
5+
[JsonPropertyName("task_detail")]
6+
public string Task { get; set; } = "";
7+
8+
//[JsonPropertyName("reason")]
9+
//public string Reason { get; set; } = "";
10+
11+
[JsonPropertyName("step")]
12+
public int Step { get; set; } = -1;
13+
14+
[JsonPropertyName("need_breakdown_task")]
15+
public bool NeedAdditionalInformation { get; set; } = false;
16+
17+
[JsonPropertyName("need_lookup_dictionary")]
18+
public bool NeedLookupDictionary { get; set; } = false;
19+
20+
[JsonPropertyName("related_tables")]
21+
public string[] Tables { get; set; } = [];
22+
23+
[JsonPropertyName("has_found_relevant_knowledge")]
24+
public bool HasFoundRelevantKnowledge { get; set; } = false;
25+
26+
//[JsonPropertyName("related_urls")]
27+
//public string[] Urls { get; set; } = [];
28+
29+
//[JsonPropertyName("input_args")]
30+
//public JsonDocument[] Parameters { get; set; } = [];
31+
32+
//[JsonPropertyName("output_results")]
33+
//public string[] Results { get; set; } = [];
34+
35+
public override string ToString()
36+
{
37+
return $"STEP {Step}: {Task}";
38+
}
39+
}

‎src/Plugins/BotSharp.Plugin.Planner/SqlGeneration/Models/PrimaryRequirementRequest.cs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
namespace BotSharp.Plugin.Planner.SqlGeneration.Models;
2+
3+
public class PrimaryRequirementRequest
4+
{
5+
[JsonPropertyName("requirement_detail")]
6+
public string Requirements { get; set; } = null!;
7+
8+
[JsonPropertyName("questions")]
9+
public string[] Questions { get; set; } = [];
10+
11+
[JsonPropertyName("norm_questions")]
12+
public string[] NormQuestions { get; set; } = [];
13+
}

‎src/Plugins/BotSharp.Plugin.Planner/SqlGeneration/Models/SecondStagePlan.cs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
namespace BotSharp.Plugin.Planner.SqlGeneration.Models;
2+
3+
public class SecondStagePlan
4+
{
5+
[JsonPropertyName("related_tables")]
6+
public string[] Tables { get; set; } = [];
7+
8+
[JsonPropertyName("need_lookup_dictionary")]
9+
public bool NeedLookupDictionary { get; set; } = false;
10+
11+
[JsonPropertyName("description")]
12+
public string Description { get; set; } = "";
13+
14+
[JsonPropertyName("input_args")]
15+
public JsonDocument[] Parameters { get; set; } = [];
16+
17+
[JsonPropertyName("output_results")]
18+
public string[] Results { get; set; } = [];
19+
}

0 commit comments

Comments
 (0)