Merge pull request #717 from Joannall/master

Add Database Hook
This commit is contained in:
Haiping 2024-10-31 15:56:43 -05:00 committed by GitHub
commit 100a31d9d4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 104 additions and 33 deletions

View file

@ -0,0 +1,7 @@
namespace BotSharp.Abstraction.Repositories;
public interface IDatabaseHook
{
// Get database type
string GetDatabaseType(RoleDialogModel message);
}

View file

@ -16,9 +16,6 @@
<None Remove="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\functions\plan_secondary_stage.json" />
<None Remove="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\functions\plan_summary.json" />
<None Remove="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\instructions\instruction.liquid" />
<None Remove="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\database.dictionary.sql.liquid" />
<None Remove="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\database.summarize.MySql.liquid" />
<None Remove="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\database.summarize.SqlServer.liquid" />
<None Remove="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\two_stage.2nd.plan.liquid" />
<None Remove="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\two_stage.next.liquid" />
<None Remove="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\two_stage.summarize.liquid" />
@ -44,15 +41,6 @@
<Content Include="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\instructions\instruction.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Include="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\database.summarize.mysql.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Include="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\database.summarize.sqlserver.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Include="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\database.dictionary.sql.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Include="data\agents\282a7128-69a1-44b0-878c-a9159b88f3b9\templates\two_stage.next.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>

View file

@ -29,6 +29,10 @@
<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\database.dictionary.sql.liquid" />
<None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\database.summarize.mysql.liquid" />
<None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\database.summarize.redshift.liquid" />
<None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\database.summarize.sqlserver.liquid" />
<None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\query_result_formatting.liquid" />
<None Remove="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\sql_statement_correctness.liquid" />
</ItemGroup>
@ -70,6 +74,18 @@
<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\database.summarize.redshift.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Include="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\database.dictionary.sql.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Include="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\database.summarize.mysql.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Include="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\database.summarize.sqlserver.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Include="data\agents\beda4c12-e1ec-4b4b-b328-3df4a6687c4f\templates\sql_statement_correctness.liquid">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
@ -81,6 +97,7 @@
<ItemGroup>
<PackageReference Include="Microsoft.AspNetCore.Mvc.Core" Version="2.2.5" />
<PackageReference Include="MySqlConnector" Version="2.3.7" />
<PackageReference Include="Npgsql" Version="8.0.5" />
</ItemGroup>
<ItemGroup>

View file

@ -1,4 +1,5 @@
using BotSharp.Abstraction.Agents.Enums;
using BotSharp.Abstraction.Repositories;
using BotSharp.Abstraction.Routing;
using BotSharp.Core.Infrastructures;
using BotSharp.Plugin.SqlDriver.Models;
@ -6,6 +7,7 @@ using Dapper;
using Microsoft.Data.SqlClient;
using Microsoft.Extensions.Logging;
using MySqlConnector;
using Npgsql;
namespace BotSharp.Plugin.SqlDriver.Functions;
@ -27,18 +29,18 @@ public class ExecuteQueryFn : IFunctionCallback
public async Task<bool> Execute(RoleDialogModel message)
{
var args = JsonSerializer.Deserialize<ExecuteQueryArgs>(message.FunctionArgs);
var refinedArgs = await RefineSqlStatement(message, args);
var settings = _services.GetRequiredService<SqlDriverSetting>();
var dbHook = _services.GetRequiredService<IDatabaseHook>();
var dbType = dbHook.GetDatabaseType(message);
try
{
var results = settings.DatabaseType switch
var results = dbType.ToLower() switch
{
"MySql" => RunQueryInMySql(refinedArgs.SqlStatements),
"SqlServer" => RunQueryInSqlServer(refinedArgs.SqlStatements),
_ => throw new NotImplementedException($"Database type {settings.DatabaseType} is not supported.")
"mysql" => RunQueryInMySql(refinedArgs.SqlStatements),
"sqlserver" => RunQueryInSqlServer(refinedArgs.SqlStatements),
"redshift" => RunQueryInRedshift(refinedArgs.SqlStatements),
_ => throw new NotImplementedException($"Database type {dbType} is not supported.")
};
if (refinedArgs.SqlStatements.Length == 1 && refinedArgs.SqlStatements[0].StartsWith("DROP TABLE"))
@ -102,6 +104,12 @@ public class ExecuteQueryFn : IFunctionCallback
using var connection = new SqlConnection(settings.SqlServerExecutionConnectionString ?? settings.SqlServerConnectionString);
return connection.Query(string.Join("\r\n", sqlTexts));
}
private IEnumerable<dynamic> RunQueryInRedshift(string[] sqlTexts)
{
var settings = _services.GetRequiredService<SqlDriverSetting>();
using var connection = new NpgsqlConnection(settings.RedshiftConnectionString);
return connection.Query(string.Join("\r\n", sqlTexts));
}
private async Task<ExecuteQueryArgs> RefineSqlStatement(RoleDialogModel message, ExecuteQueryArgs args)
{

View file

@ -1,7 +1,12 @@
using BotSharp.Abstraction.Repositories;
using BotSharp.Plugin.SqlDriver.Models;
using Microsoft.Data.SqlClient;
using Microsoft.EntityFrameworkCore.Metadata.Internal;
using Microsoft.Extensions.Logging;
using MySqlConnector;
using Npgsql;
using System.ComponentModel.DataAnnotations;
using System.Data.Common;
namespace BotSharp.Plugin.SqlDriver.Functions;
@ -25,14 +30,16 @@ public class GetTableDefinitionFn : IFunctionCallback
var args = JsonSerializer.Deserialize<SqlStatement>(message.FunctionArgs);
var tables = args.Tables;
var agentService = _services.GetRequiredService<IAgentService>();
var settings = _services.GetRequiredService<SqlDriverSetting>();
var dbHook = _services.GetRequiredService<IDatabaseHook>();
var dbType = dbHook.GetDatabaseType(message);
// Get table DDL from database
var tableDdls = settings.DatabaseType switch
var tableDdls = dbType switch
{
"MySql" => GetDdlFromMySql(tables),
"SqlServer" => GetDdlFromSqlServer(tables),
_ => throw new NotImplementedException($"Database type {settings.DatabaseType} is not supported.")
"mysql" => GetDdlFromMySql(tables),
"sqlserver" => GetDdlFromSqlServer(tables),
"redshift" => GetDdlFromRedshift(tables),
_ => throw new NotImplementedException($"Database type {dbType} is not supported.")
};
message.Content = string.Join("\r\n\r\n", tableDdls);
@ -125,4 +132,43 @@ public class GetTableDefinitionFn : IFunctionCallback
connection.Close();
return tableDdls;
}
private List<string> GetDdlFromRedshift(string[] tables)
{
var settings = _services.GetRequiredService<SqlDriverSetting>();
var tableDdls = new List<string>();
var schemas = "'onebi_hour','onebi_day'";
using var connection = new NpgsqlConnection(settings.RedshiftConnectionString);
connection.Open();
foreach (var table in tables)
{
try
{
var sql = $"SELECT columnname, external_type, columnnum FROM svv_external_columns WHERE schemaname in ({schemas}) AND tablename = '{table}';";
using var command = new NpgsqlCommand(sql, connection);
using var reader = command.ExecuteReader();
var ddl = $"Table {table} contains below fields: \r\n";
while (reader.Read())
{
/*var result = reader.GetString(0);
tableDdls.Add(result);*/
var columnName = reader.GetString(0);
var dataType = reader.GetString(1);
ddl += $"{columnName} {dataType}\n";
}
tableDdls.Add(ddl);
reader.Close();
command.Dispose();
}
catch (Exception ex)
{
_logger.LogWarning($"Error when getting ddl statement of table {table}. {ex.Message}\r\n{ex.InnerException}");
}
}
connection.Close();
return tableDdls;
}
}

View file

@ -14,14 +14,6 @@ public class SqlDriverPlanningHook : IPlanningHook
_services = services;
}
public async Task<string> GetSummaryAdditionalRequirements(string planner)
{
var settings = _services.GetRequiredService<SqlDriverSetting>();
var agentService = _services.GetRequiredService<IAgentService>();
var agent = await agentService.GetAgent(BuiltInAgentId.Planner);
return agent.Templates.FirstOrDefault(x => x.Name == $"database.summarize.{settings.DatabaseType.ToLower()}")?.Content ?? string.Empty;
}
public async Task OnPlanningCompleted(string planner, RoleDialogModel msg)
{
var settings = _services.GetRequiredService<SqlDriverSetting>();

View file

@ -2,7 +2,7 @@ namespace BotSharp.Plugin.SqlHero.Settings;
public class SqlDriverSetting
{
public string DatabaseType { get; set; } = "MySql";
public string DatabaseType { get; set; } = "mysql";
public string MySqlConnectionString { get; set; } = null!;
public string MySqlExecutionConnectionString { get; set; } = null!;
public string MySqlTempConnectionString { get; set; } = null!;
@ -10,6 +10,7 @@ public class SqlDriverSetting
public string SqlServerConnectionString { get; set; } = null!;
public string SqlServerExecutionConnectionString { get; set; } = null!;
public string SqlLiteConnectionString { get; set; } = null!;
public string RedshiftConnectionString { get; set; } = null!;
public bool ExecuteSqlSelectAutonomous { get; set; } = false;
public bool FormattingResult { get; set; } = true;
}

View file

@ -0,0 +1,12 @@
Try if you can generate a Redshift single query to fulfill the needs. The step should contains all needed parameters.
The parameters can be extracted from the original task.
If not, generate the query step by step based on the planning.
The query must exactly based on the provided table structure. And carefully review the foreign keys to make sure you include all the accurate information.
Note: Output should be only the sql query with sql comments that can be directly run in Amazon Redshift.
*** the alias of the table name in the sql query should be identical. ***
*** The generated sql query MUST be based on the provided table structure. ***
*** All queries return a maximum of 1000 records. ***
*** Only select user friendly columns. ***