feat(ai): support SQL library file references

This commit is contained in:
zipg
2026-07-07 15:39:10 +08:00
committed by GitHub
parent fc01cea80f
commit d90847bbba
15 changed files with 667 additions and 178 deletions
File diff suppressed because it is too large Load Diff
+3 -2
View File
@@ -1205,8 +1205,9 @@ export default {
thinking: "Thinking...",
reasoningProcess: "Thinking process",
stopGenerating: "Stop generating",
tableMentionEmpty: "No matching tables",
tableMentionPlaceholderHint: "Type {'@'} to mention tables for context",
scrollToBottom: "Back to bottom",
tableMentionEmpty: "No matching tables or SQL files",
tableMentionPlaceholderHint: "Type {'@'} to mention tables or SQL files for context",
tableMentionTypes: {
table: "Table",
view: "View",
+3 -2
View File
@@ -1178,8 +1178,9 @@ export default withEnglishFallback({
thinking: "Procesando...",
reasoningProcess: "Proceso de razonamiento",
stopGenerating: "Detener generación",
tableMentionEmpty: "Sin tablas coincidentes",
tableMentionPlaceholderHint: "Escribe {'@'} para mencionar tablas como contexto",
scrollToBottom: "Volver al final",
tableMentionEmpty: "Sin tablas ni archivos SQL coincidentes",
tableMentionPlaceholderHint: "Escribe {'@'} para mencionar tablas o archivos SQL como contexto",
tableMentionTypes: {
table: "Tabla",
view: "Vista",
+3 -2
View File
@@ -1159,8 +1159,9 @@ export default withEnglishFallback({
thinking: "Riflessione...",
reasoningProcess: "Processo di riflessione",
stopGenerating: "Interrompi generazione",
tableMentionEmpty: "Nessuna tabella corrispondente",
tableMentionPlaceholderHint: "Digita {'@'} per menzionare le tabelle per il contesto",
scrollToBottom: "Torna in fondo",
tableMentionEmpty: "Nessuna tabella o file SQL corrispondente",
tableMentionPlaceholderHint: "Digita {'@'} per menzionare tabelle o file SQL per il contesto",
tableMentionTypes: {
table: "Tabella",
view: "Vista",
+3 -2
View File
@@ -1181,8 +1181,9 @@ export default withEnglishFallback({
thinking: "思考中...",
reasoningProcess: "思考プロセス",
stopGenerating: "生成を停止",
tableMentionEmpty: "一致するテーブルがありません",
tableMentionPlaceholderHint: "{'@'}を入力してコンテキスト用のテーブルを指定",
scrollToBottom: "一番下へ戻る",
tableMentionEmpty: "一致するテーブルまたは SQL ファイルがありません",
tableMentionPlaceholderHint: "{'@'}を入力してコンテキスト用のテーブルまたは SQL ファイルを指定",
tableMentionTypes: {
table: "テーブル",
view: "ビュー",
+3 -2
View File
@@ -1179,8 +1179,9 @@ export default withEnglishFallback({
thinking: "Pensando...",
reasoningProcess: "Processo de raciocínio",
stopGenerating: "Parar geração",
tableMentionEmpty: "Nenhuma tabela correspondente",
tableMentionPlaceholderHint: "Digite {'@'} para mencionar tabelas como contexto",
scrollToBottom: "Voltar ao fim",
tableMentionEmpty: "Nenhuma tabela ou arquivo SQL correspondente",
tableMentionPlaceholderHint: "Digite {'@'} para mencionar tabelas ou arquivos SQL como contexto",
tableMentionTypes: {
table: "Tabela",
view: "Visão",
+3 -2
View File
@@ -1207,8 +1207,9 @@ export default withEnglishFallback({
thinking: "思考中...",
reasoningProcess: "思考过程",
stopGenerating: "停止生成",
tableMentionEmpty: "没有匹配的表",
tableMentionPlaceholderHint: "输入 {'@'} 可选择表并带入上下文",
scrollToBottom: "回到底部",
tableMentionEmpty: "没有匹配的表或 SQL 文件",
tableMentionPlaceholderHint: "输入 {'@'} 可选择表或 SQL 文件并带入上下文",
tableMentionTypes: {
table: "表",
view: "视图",
+3 -2
View File
@@ -1160,8 +1160,9 @@ export default withEnglishFallback({
thinking: "思考中……",
reasoningProcess: "思考過程",
stopGenerating: "停止產生",
tableMentionEmpty: "沒有相符的資料表",
tableMentionPlaceholderHint: "輸入 {'@'} 可提及資料表作為脈絡",
scrollToBottom: "回到底部",
tableMentionEmpty: "沒有相符的資料表或 SQL 檔案",
tableMentionPlaceholderHint: "輸入 {'@'} 可提及資料表或 SQL 檔案作為脈絡",
tableMentionTypes: {
table: "資料表",
view: "檢視",
@@ -10,6 +10,7 @@ function context(overrides: Partial<AiContext> = {}): AiContext {
database: "app",
currentSql: "",
tables: [],
sqlFiles: [],
truncated: false,
...overrides,
};
+28 -1
View File
@@ -66,6 +66,13 @@ export interface AiSchemaTable {
foreignKeys?: ForeignKeyInfo[];
}
export interface AiSqlFileContext {
id: string;
name: string;
sql: string;
truncated?: boolean;
}
export interface AiContext {
connectionId: string;
connectionName: string;
@@ -75,6 +82,7 @@ export interface AiContext {
lastError?: string;
lastResultPreview?: string;
tables: AiSchemaTable[];
sqlFiles: AiSqlFileContext[];
schemaScope?: "focused_table" | "database";
truncated: boolean;
}
@@ -196,6 +204,7 @@ export function buildSystemPrompt(action: AiAction, context: AiContext, mode: Ai
const schema = formatSchema(context);
const resultPreview = context.lastResultPreview ? `\nLast result preview:\n${context.lastResultPreview}\n` : "";
const lastError = context.lastError ? `\nLast error:\n${context.lastError}\n` : "";
const referencedSqlFiles = formatReferencedSqlFiles(context);
const schemaScope = context.schemaScope ?? "database";
const isZh = isChineseLocale(currentLocale());
@@ -225,6 +234,7 @@ export function buildSystemPrompt(action: AiAction, context: AiContext, mode: Ai
schemaCoverageLine(context, isZh),
"",
`Current SQL:\n${context.currentSql.trim() || "(empty)"}`,
referencedSqlFiles,
lastError,
resultPreview,
`Schema:\n${schema}`,
@@ -251,6 +261,7 @@ function buildBasePromptLines(isZh: boolean): string[] {
isZh ? "当用户要求分析或查看某个表时,生成 SELECT 查询获取数据,而不是查询元数据。" : "When the user asks to 'analyze' or 'look at' a table, generate a SELECT query to retrieve data, not a metadata query.",
isZh ? "不要编造 Schema 中不存在的表或列。" : "Never invent tables or columns that are not in the schema context.",
isZh ? "用户输入中的 @schema.table 或 @table 表示用户明确提到的表;这些表已优先放入 Schema 上下文。" : "User input may contain @schema.table or @table mentions. Treat them as explicit table references; mentioned tables are prioritized in the schema context.",
isZh ? "用户引用的 SQL 库文件是额外上下文;可参考其中的查询意图、业务过滤条件和 SQL 写法,但不要把它当作当前编辑器内容。" : "Referenced SQL library files are additional context. Use them for query intent, business filters, and SQL style, but do not treat them as the current editor content.",
isZh ? "不要生成多语句 SQL,除非用户明确要求。不要在同一个回答里混合 SELECT 和写操作。" : "Do not generate multi-statement SQL unless the user explicitly asks for it. Do not mix SELECT statements and write operations in the same answer.",
isZh ? "对于 DROP、DELETE、TRUNCATE、ALTER 或没有 WHERE 的 UPDATE,简要警告并优先提供安全的 SELECT 预览。" : "For destructive statements (DROP, DELETE, TRUNCATE, ALTER, UPDATE without WHERE), warn briefly and prefer a safer SELECT preview.",
isZh ? "对于 UPDATE 或 DELETE,必须带 WHERE 并说明影响范围;生产库写操作只给建议,不主动建议执行。" : "For UPDATE or DELETE, require a WHERE clause and explain the affected scope; for production writes, provide guidance but do not proactively suggest execution.",
@@ -263,6 +274,7 @@ function buildVectorSystemPrompt(context: AiContext, mode: AiAssistantMode): str
const schema = formatSchema(context);
const resultPreview = context.lastResultPreview ? `\nLast result preview:\n${context.lastResultPreview}\n` : "";
const lastError = context.lastError ? `\nLast error:\n${context.lastError}\n` : "";
const referencedSqlFiles = formatReferencedSqlFiles(context);
const lines: string[] = [
isZh ? `你是 DBX 内置的向量数据库助手。当前连接的是 ${dbLabel(context.databaseType)} 数据库。用中文回复。` : `You are DBX's vector database assistant. Connected to ${dbLabel(context.databaseType)}. Reply in English.`,
isZh ? "数据存储在集合(collections)中,每条记录包含唯一标识及可选的元数据负载(payload/metadata)。" : "Data is stored in collections. Each record has a unique identifier and optional metadata payload.",
@@ -274,6 +286,7 @@ function buildVectorSystemPrompt(context: AiContext, mode: AiAssistantMode): str
schemaCoverageLine(context, isZh),
"",
`Current collection:\n${context.currentSql.trim() || "(none)"}`,
referencedSqlFiles,
lastError,
resultPreview,
"",
@@ -369,7 +382,20 @@ function formatSchema(context: AiContext): string {
.join("\n\n");
}
export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig, options: { maxTables?: number; maxColumnsPerTable?: number; maxIndexesPerTable?: number; maxFksPerTable?: number; mentionedTables?: AiTableMention[] } = {}): Promise<AiContext> {
function formatReferencedSqlFiles(context: AiContext): string {
if (!context.sqlFiles.length) return "";
return [
"Referenced SQL library files:",
...context.sqlFiles.map((file) => {
const sql = file.sql.trim() || "(empty)";
const suffix = file.truncated ? " (truncated)" : "";
return `File: ${file.name}${suffix}\nSQL:\n${sql}`;
}),
].join("\n\n");
}
export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig, options: { maxTables?: number; maxColumnsPerTable?: number; maxIndexesPerTable?: number; maxFksPerTable?: number; mentionedTables?: AiTableMention[]; sqlFiles?: AiSqlFileContext[] } = {}): Promise<AiContext> {
const maxTables = options.maxTables ?? 50;
const maxColumnsPerTable = options.maxColumnsPerTable ?? 40;
const maxIndexesPerTable = options.maxIndexesPerTable ?? 10;
@@ -496,6 +522,7 @@ export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig
lastError: extractLastError(tab.result),
lastResultPreview: formatResultPreview(tab.result),
tables,
sqlFiles: options.sqlFiles ?? [],
schemaScope,
truncated,
};
+1
View File
@@ -538,6 +538,7 @@ export async function listSqlFilesInFolder(folderPath: string): Promise<SqlFileE
export interface AiChatMessage {
role: string;
content: string;
mentions?: unknown[];
reasoning?: string;
kind?: "contextSummary";
}
+2
View File
@@ -194,6 +194,8 @@ pub struct AiChatMessage {
pub role: String,
pub content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mentions: Option<serde_json::Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub kind: Option<String>,
+25
View File
@@ -34,6 +34,7 @@ const { buildSystemPrompt, isVectorDbType, buildUserPrompt } = await import("../
function context(overrides: Partial<AiContext> = {}): AiContext {
return {
connectionId: "conn-1",
connectionName: "prod-analytics",
databaseType: "postgres",
database: "app",
@@ -52,6 +53,7 @@ function context(overrides: Partial<AiContext> = {}): AiContext {
foreignKeys: [{ column: "user_id", ref_table: "users", ref_column: "id" }],
},
],
sqlFiles: [],
truncated: false,
...overrides,
};
@@ -99,10 +101,32 @@ test("prompt enforces database dialect and single executable statement safety",
assert.match(prompt, /不要在同一个回答里混合 SELECT 和写操作/);
});
test("prompt includes referenced SQL library files as additional context", () => {
const prompt = buildSystemPrompt(
"generate",
context({
sqlFiles: [
{
id: "sql-1",
name: "monthly-orders.sql",
sql: "select date_trunc('month', created_at) as month, count(*) from public.orders group by 1",
},
],
}),
"ask",
);
assert.match(prompt, /SQL 库文件是额外上下文/);
assert.match(prompt, /Referenced SQL library files:/);
assert.match(prompt, /File: monthly-orders\.sql/);
assert.match(prompt, /date_trunc\('month', created_at\)/);
});
// Vector database tests
function vectorContext(overrides: Partial<AiContext> = {}): AiContext {
return {
connectionId: "conn-1",
connectionName: "my-qdrant",
databaseType: "qdrant",
database: "default",
@@ -115,6 +139,7 @@ function vectorContext(overrides: Partial<AiContext> = {}): AiContext {
columns: [],
},
],
sqlFiles: [],
truncated: false,
...overrides,
};
+2
View File
@@ -34,6 +34,7 @@ const { buildSystemPrompt } = await import("../../apps/desktop/src/lib/ai/ai.ts"
function baseContext(overrides: Partial<AiContext> = {}): AiContext {
return {
connectionId: "conn-1",
connectionName: "prod-analytics",
databaseType: "postgres",
database: "warehouse",
@@ -67,6 +68,7 @@ function baseContext(overrides: Partial<AiContext> = {}): AiContext {
],
},
],
sqlFiles: [],
truncated: false,
...overrides,
};
@@ -4,6 +4,7 @@ import { test } from "vitest";
import { buildSystemPrompt, type AiContext } from "../../apps/desktop/src/lib/ai/ai.ts";
const BASE_CONTEXT: AiContext = {
connectionId: "conn-1",
connectionName: "local",
databaseType: "mysql",
database: "app",
@@ -35,6 +36,7 @@ const BASE_CONTEXT: AiContext = {
],
},
],
sqlFiles: [],
truncated: false,
};