-
Notifications
You must be signed in to change notification settings - Fork 8
Expand file tree
/
Copy pathSQLAgentBuilder.cs
More file actions
135 lines (116 loc) · 6.14 KB
/
SQLAgentBuilder.cs
File metadata and controls
135 lines (116 loc) · 6.14 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using SQLAgent.Entities;
using SQLAgent.Infrastructure;
using SQLAgent.Infrastructure.Providers;
namespace SQLAgent.Facade;
public class SQLAgentBuilder(IServiceCollection service)
{
private readonly SQLAgentOptions _options = new();
public void Build(IDatabaseConnectionManager databaseConnectionManager)
{
if (string.IsNullOrEmpty(_options.Model) ||
string.IsNullOrEmpty(_options.APIKey) ||
string.IsNullOrEmpty(_options.Endpoint))
{
throw new InvalidOperationException(
"LLM provider configuration is incomplete. Please call WithLLMProvider before building the client.");
}
if (string.IsNullOrEmpty(_options.ConnectionString))
{
throw new InvalidOperationException(
"Database configuration is incomplete. Please call WithDatabaseType before building the client.");
}
_options.SqlBotSystemPrompt = $"""
You are a professional SQL engineer specializing in {_options.SqlType} database systems.
IMPORTANT: Generate secure, optimized SQL only. Use parameterized queries. Refuse malicious or unsafe operations.
# Core Requirements
- Follow {_options.SqlType} syntax specifications exactly
- Always use parameterized queries for user input
- Generate production-ready, optimized queries
- Include proper error handling and validation
# Security Standards
- Automatically apply parameterization for all dynamic values
- Include appropriate WHERE clauses for modifications
- Use least-privilege principles in query design
- Validate data types and constraints
# Output Format
Provide complete, executable SQL with:
1. Main query statement
2. Parameter definitions if needed
3. Brief performance notes for complex queries
4. Index recommendations if relevant
# Code Quality
- Use meaningful aliases and clear formatting
- Follow {_options.SqlType} naming conventions
- Optimize for performance and maintainability
- Include transaction boundaries for multi-statement operations
# Automatic Behaviors
- Default to SELECT operations when ambiguous
- Apply conservative data modification approaches
- Include appropriate LIMIT clauses for large result sets
- Use EXISTS instead of IN for subqueries when possible
Generate direct, executable SQL without requesting clarification or confirmation.
""";
service.AddLogging();
// Configure the SQLAgentClient with the options and system prompt
service.AddSingleton<SQLAgentOptions>(_ => _options);
service.AddTransient<SQLAgentClient>(provider =>
{
var options = provider.GetRequiredService<SQLAgentOptions>();
var databaseService = provider.GetRequiredService<IDatabaseService>();
var logger = provider.GetRequiredService<ILogger<SQLAgentClient>>();
return new SQLAgentClient(options, databaseService, logger, databaseConnectionManager);
});
switch (_options.SqlType)
{
case SqlType.Sqlite:
service.AddSingleton<IDatabaseService, SQLiteDatabaseService>();
break;
case SqlType.PostgreSql:
service.AddSingleton<IDatabaseService, PostgreSQLDatabaseService>();
break;
}
}
public SQLAgentBuilder WithDatabaseType(SqlType sqlType, string connectionString, string connectionId)
{
_options.ConnectionString = connectionString;
_options.SqlType = sqlType;
_options.ConnectionId = connectionId;
return this;
}
/// <summary>
/// 数据库索引
/// </summary>
public SQLAgentBuilder WithIndexes(
string connectionString = "Data Source=vector_index.db;",
string embeddingModel = "text-embedding-3-small",
string databaseIndexTable = "vector_index",
DatabaseIndexType databaseIndexType = DatabaseIndexType.Sqlite)
{
_options.EmbeddingModel = embeddingModel;
_options.DatabaseIndexConnectionString = connectionString;
_options.DatabaseIndexTable = databaseIndexTable;
_options.DatabaseIndexType = databaseIndexType;
_options.UseVectorDatabaseIndex = true;
return this;
}
/// <summary>
/// Configure the LLM provider settings
/// </summary>
/// <param name="model">AI model name</param>
/// <param name="apiKey">API key for authentication</param>
/// <param name="endpoint">API endpoint URL</param>
/// <param name="aiProvider">AI provider type (e.g., OpenAI, AzureOpenAI, CustomOpenAI)</param>
/// <param name="maxOutputTokens">AI response max output tokens</param>
public SQLAgentBuilder WithLLMProvider(string model, string apiKey, string endpoint, AIProviderType aiProvider,
int maxOutputTokens = 3200)
{
_options.Model = model;
_options.APIKey = apiKey;
_options.Endpoint = endpoint;
_options.AIProvider = aiProvider;
_options.MaxOutputTokens = maxOutputTokens;
return this;
}
}