feat: enhance AI functionality with conversation management and cancellation

src/lib/ai.ts

@@ -1,5 +1,5 @@
 import type { AiConfig } from "@/stores/settingsStore";
-import type { ColumnInfo, ConnectionConfig, DatabaseType, QueryResult, QueryTab } from "@/types/database";
+import type { ColumnInfo, ConnectionConfig, DatabaseType, ForeignKeyInfo, IndexInfo, QueryResult, QueryTab } from "@/types/database";
 import * as api from "@/lib/tauri";
 import { currentLocale } from "@/i18n";
 
@@ -10,6 +10,8 @@ export interface AiSchemaTable {
   name: string;
   tableType: string;
   columns: ColumnInfo[];
+  indexes?: IndexInfo[];
+  foreignKeys?: ForeignKeyInfo[];
 }
 
 export interface AiContext {
@@ -30,20 +32,40 @@ export interface AiRequestInput {
   context: AiContext;
 }
 
-const ACTION_INSTRUCTIONS: Record<AiAction, string> = {
-  generate: "Generate a SQL query that satisfies the user's request.",
-  explain: "Explain the current SQL clearly and point out risky operations or assumptions.",
-  optimize: "Rewrite or suggest improvements for the current SQL. Prefer a complete improved SQL query first, followed by short notes.",
-  fix: "Fix the current SQL using the provided error/result context. Return the corrected SQL first, followed by short notes if needed.",
-  convert: "Convert the current SQL to the target dialect requested by the user. Return the converted SQL first.",
-  sampleData: "Generate safe sample SQL statements or mock data for the current schema. Do not use real production data.",
+const ACTION_INSTRUCTIONS: Record<AiAction, { en: string; zh: string }> = {
+  generate: {
+    en: "Generate a SQL query that satisfies the user's request. Return the SQL in a ```sql code block first, followed by a brief note if needed. Use foreign key relationships from the schema to infer correct JOIN conditions.",
+    zh: "根据用户请求生成 SQL。先在 ```sql 代码块中返回 SQL，必要时附简短说明。利用 Schema 中的外键关系推断正确的 JOIN 条件。",
+  },
+  explain: {
+    en: "Explain the current SQL step by step. Point out risky operations, implicit assumptions, and potential performance issues. Reference index and foreign key info from the schema when relevant.",
+    zh: "逐步解释当前 SQL。指出危险操作、隐含假设和潜在性能问题。结合 Schema 中的索引和外键信息分析。",
+  },
+  optimize: {
+    en: "Rewrite or suggest improvements for the current SQL. Return the improved SQL in a ```sql code block first, followed by short notes explaining the changes. Use the index information in the schema to suggest index-aware optimizations (e.g., avoid full table scans, leverage existing indexes).",
+    zh: "重写或优化当前 SQL。先在 ```sql 代码块中返回优化后的 SQL，然后简要说明改动。利用 Schema 中的索引信息建议索引友好的优化（如避免全表扫描、利用现有索引）。",
+  },
+  fix: {
+    en: "Fix the current SQL using the provided error message and result context. Return the corrected SQL in a ```sql code block first, followed by a brief explanation of the root cause.",
+    zh: "根据报错信息和结果上下文修复当前 SQL。先在 ```sql 代码块中返回修正后的 SQL，再简要说明根因。",
+  },
+  convert: {
+    en: "Convert the current SQL to the target dialect requested by the user. Return the converted SQL in a ```sql code block first. Note any syntax differences or incompatibilities.",
+    zh: "将当前 SQL 转换为用户指定的目标方言。先在 ```sql 代码块中返回转换后的 SQL，再说明语法差异。",
+  },
+  sampleData: {
+    en: "Generate safe sample INSERT statements or mock data for the current schema. Do not use real production data. Return SQL in a ```sql code block.",
+    zh: "为当前 Schema 生成安全的示例 INSERT 语句或模拟数据。不使用真实生产数据。在 ```sql 代码块中返回 SQL。",
+  },
 };
 
 export async function runAiAction(input: AiRequestInput, history?: api.AiMessage[]): Promise<string> {
+  const isZh = currentLocale() === "zh-CN";
   const systemPrompt = buildSystemPrompt(input.action, input.context);
+  const instruction = isZh ? ACTION_INSTRUCTIONS[input.action].zh : ACTION_INSTRUCTIONS[input.action].en;
   const userPrompt = [
     `Action: ${input.action}`,
-    ACTION_INSTRUCTIONS[input.action],
+    instruction,
     "",
     "User request:",
     input.instruction.trim() || "(No extra instruction provided.)",
@@ -54,12 +76,13 @@ export async function runAiAction(input: AiRequestInput, history?: api.AiMessage
     { role: "user", content: userPrompt },
   ];
 
+  const params = actionParams(input.action);
   return api.aiComplete({
     config: input.config,
     systemPrompt,
     messages,
-    maxTokens: 2400,
-    temperature: 0.15,
+    maxTokens: params.maxTokens,
+    temperature: params.temperature,
   });
 }
 
@@ -67,11 +90,14 @@ export async function runAiStream(
   input: AiRequestInput,
   history: api.AiMessage[] | undefined,
   onDelta: (delta: string) => void,
+  sessionId?: string,
 ): Promise<void> {
+  const isZh = currentLocale() === "zh-CN";
   const systemPrompt = buildSystemPrompt(input.action, input.context);
+  const instruction = isZh ? ACTION_INSTRUCTIONS[input.action].zh : ACTION_INSTRUCTIONS[input.action].en;
   const userPrompt = [
     `Action: ${input.action}`,
-    ACTION_INSTRUCTIONS[input.action],
+    instruction,
     "",
     "User request:",
     input.instruction.trim() || "(No extra instruction provided.)",
@@ -82,19 +108,28 @@ export async function runAiStream(
     { role: "user", content: userPrompt },
   ];
 
-  const sessionId = crypto.randomUUID();
+  const sid = sessionId || crypto.randomUUID();
+  const params = actionParams(input.action);
 
-  await api.aiStream(sessionId, {
+  await api.aiStream(sid, {
     config: input.config,
     systemPrompt,
     messages,
-    maxTokens: 2400,
-    temperature: 0.15,
+    maxTokens: params.maxTokens,
+    temperature: params.temperature,
   }, (chunk) => {
     if (!chunk.done && chunk.delta) onDelta(chunk.delta);
   });
 }
 
+function actionParams(action: AiAction): { maxTokens: number; temperature: number } {
+  switch (action) {
+    case "explain": return { maxTokens: 3200, temperature: 0.2 };
+    case "sampleData": return { maxTokens: 2400, temperature: 0.1 };
+    default: return { maxTokens: 2400, temperature: 0.15 };
+  }
+}
+
 export function extractSql(text: string): string {
   const fenced = text.match(/```(?:sql|mysql|postgresql|sqlite|tsql|clickhouse)?\s*([\s\S]*?)```/i);
   if (fenced?.[1]) return fenced[1].trim();
@@ -110,7 +145,7 @@ export function buildSystemPrompt(action: AiAction, context: AiContext): string
 
   const isZh = currentLocale() === "zh-CN";
 
-  return [
+  const lines: string[] = [
     isZh
       ? "你是 DBX 内置的数据库助手。用中文回复。"
       : "You are DBX's built-in database assistant. Reply in English.",
@@ -118,34 +153,50 @@ export function buildSystemPrompt(action: AiAction, context: AiContext): string
       ? "精确、保守，根据当前数据库方言生成 SQL。"
       : "Be precise, conservative, and adapt SQL to the active database dialect.",
     isZh
-      ? "下面的 Schema 上下文已包含表和列信息，直接使用即可。不要查询 information_schema 或系统表来获取结构信息，直接针对用户的实际表编写查询。"
-      : "The schema context below already contains table and column information — use it directly. Do NOT query information_schema or system tables to discover schema; write queries against the user's actual tables.",
+      ? "下面的 Schema 上下文已包含表、列、索引和外键信息，直接使用即可。不要查询 information_schema 或系统表来获取结构信息。"
+      : "The schema context below already contains tables, columns, indexes, and foreign keys — use it directly. Do NOT query information_schema or system tables.",
     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 present in the schema context unless the user explicitly asks for hypothetical examples.",
+      ? "不要编造 Schema 中不存在的表或列。"
+      : "Never invent tables or columns that are not in the schema context.",
     isZh
-      ? "对于 DROP、DELETE、TRUNCATE、ALTER 或没有 WHERE 子句的 UPDATE 等危险语句，简要警告并优先提供安全的 SELECT 预览。"
-      : "For destructive statements such as DROP, DELETE, TRUNCATE, ALTER, or UPDATE without a clear WHERE clause, warn briefly and prefer a safer SELECT preview when appropriate.",
+      ? "对于 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.",
+  ];
+
+  if (action === "optimize") {
+    lines.push(isZh
+      ? "利用 Schema 中的索引信息建议优化。指出哪些查询条件可以命中索引、哪些会导致全表扫描。"
+      : "Use the index information in the schema to suggest optimizations. Point out which conditions hit indexes and which cause full table scans.");
+  } else if (action === "generate") {
+    lines.push(isZh
+      ? "利用外键关系推断 JOIN 条件。生成操作优先返回 SQL，避免长篇解释。"
+      : "Use foreign key relationships to infer JOIN conditions. Return the SQL first and avoid long explanations.");
+  } else if (action === "fix") {
+    lines.push(isZh
+      ? "仔细分析错误信息，定位根因。先返回修正后的 SQL，再简要解释。"
+      : "Carefully analyze the error message to identify the root cause. Return the corrected SQL first, then briefly explain.");
+  }
+
+  lines.push(
     isZh
-      ? "返回 SQL 时放在 ```sql 代码块中。额外说明简短实用即可。"
-      : "When returning SQL, put the SQL in a fenced ```sql code block. Keep extra explanation short and practical.",
-    action === "generate"
-      ? (isZh ? "生成操作优先返回 SQL，避免长篇解释。" : "For generate actions, return the SQL first and avoid long explanations.")
-      : "",
+      ? "返回 SQL 时放在 ```sql 代码块中。额外说明简短实用。"
+      : "Put SQL in a fenced ```sql code block. Keep extra explanation short and practical.",
     "",
     `Database type: ${context.databaseType}`,
     `Connection: ${context.connectionName}`,
     `Database: ${context.database}`,
-    context.truncated ? "Schema context is truncated." : "Schema context is complete within the current budget.",
+    context.truncated ? "Schema context is truncated." : "Schema context is complete.",
     "",
     `Current SQL:\n${context.currentSql.trim() || "(empty)"}`,
     lastError,
     resultPreview,
     `Schema:\n${schema}`,
-  ].filter(Boolean).join("\n");
+  );
+
+  return lines.filter(Boolean).join("\n");
 }
 
 function formatSchema(context: AiContext): string {
@@ -153,16 +204,33 @@ function formatSchema(context: AiContext): string {
 
   return context.tables.map((table) => {
     const name = table.schema ? `${table.schema}.${table.name}` : table.name;
-    const columns = table.columns.map((column) => {
+    const lines: string[] = [`${name} (${table.tableType})`];
+
+    for (const column of table.columns) {
       const flags = [
-        column.is_primary_key ? "primary key" : "",
-        column.is_nullable ? "nullable" : "not null",
+        column.is_primary_key ? "PK" : "",
+        column.is_nullable ? "nullable" : "NOT NULL",
         column.column_default ? `default ${column.column_default}` : "",
         column.extra || "",
       ].filter(Boolean).join(", ");
-      return `  - ${column.name}: ${column.data_type}${flags ? ` (${flags})` : ""}`;
-    });
-    return [`${name} (${table.tableType})`, ...columns].join("\n");
+      lines.push(`  - ${column.name}: ${column.data_type}${flags ? ` (${flags})` : ""}`);
+    }
+
+    if (table.indexes?.length) {
+      for (const idx of table.indexes) {
+        if (idx.is_primary) continue;
+        const unique = idx.is_unique ? "UNIQUE " : "";
+        lines.push(`  Index: ${unique}${idx.name}(${idx.columns.join(", ")})`);
+      }
+    }
+
+    if (table.foreignKeys?.length) {
+      for (const fk of table.foreignKeys) {
+        lines.push(`  FK: ${fk.column} → ${fk.ref_table}.${fk.ref_column}`);
+      }
+    }
+
+    return lines.join("\n");
   }).join("\n\n");
 }
 
@@ -177,11 +245,19 @@ export async function buildAiContext(
   let truncated = false;
 
   if (tab.tableMeta) {
+    const s = tab.tableMeta.schema ?? "";
+    const tName = tab.tableMeta.tableName;
+    const [indexes, foreignKeys] = await Promise.all([
+      api.listIndexes(tab.connectionId, tab.database, s, tName).catch(() => [] as IndexInfo[]),
+      api.listForeignKeys(tab.connectionId, tab.database, s, tName).catch(() => [] as ForeignKeyInfo[]),
+    ]);
     tables.push({
       schema: tab.tableMeta.schema,
-      name: tab.tableMeta.tableName,
+      name: tName,
       tableType: "TABLE",
       columns: tab.tableMeta.columns.slice(0, maxColumnsPerTable),
+      indexes,
+      foreignKeys,
     });
     truncated = tab.tableMeta.columns.length > maxColumnsPerTable;
   } else if (!["redis", "mongodb"].includes(connection.db_type)) {
@@ -194,12 +270,18 @@ export async function buildAiContext(
             truncated = true;
             break;
           }
-          const columns = await api.getColumns(tab.connectionId, tab.database, schema, table.name);
+          const [columns, indexes, foreignKeys] = await Promise.all([
+            api.getColumns(tab.connectionId, tab.database, schema, table.name),
+            api.listIndexes(tab.connectionId, tab.database, schema, table.name).catch(() => [] as IndexInfo[]),
+            api.listForeignKeys(tab.connectionId, tab.database, schema, table.name).catch(() => [] as ForeignKeyInfo[]),
+          ]);
           tables.push({
             schema: schema === tab.database && connection.db_type !== "postgres" ? undefined : schema,
             name: table.name,
             tableType: table.table_type,
             columns: columns.slice(0, maxColumnsPerTable),
+            indexes,
+            foreignKeys,
           });
           if (columns.length > maxColumnsPerTable) truncated = true;
         }

src/lib/tauri.ts

@@ -62,10 +62,43 @@ export async function aiTestConnection(config: AiConfig): Promise<string> {
   return invoke("ai_test_connection", { config });
 }
 
+export async function aiCancelStream(sessionId: string): Promise<boolean> {
+  return invoke("ai_cancel_stream", { sessionId });
+}
+
 export async function loadAiConfig(): Promise<AiConfig | null> {
   return invoke("load_ai_config");
 }
 
+// --- AI Conversations ---
+
+export interface AiChatMessage {
+  role: string;
+  content: string;
+}
+
+export interface AiConversation {
+  id: string;
+  title: string;
+  connectionName: string;
+  database: string;
+  messages: AiChatMessage[];
+  createdAt: string;
+  updatedAt: string;
+}
+
+export async function saveAiConversation(conversation: AiConversation): Promise<void> {
+  return invoke("save_ai_conversation", { conversation });
+}
+
+export async function loadAiConversations(): Promise<AiConversation[]> {
+  return invoke("load_ai_conversations");
+}
+
+export async function deleteAiConversation(id: string): Promise<void> {
+  return invoke("delete_ai_conversation", { id });
+}
+
 export async function testConnection(config: ConnectionConfig): Promise<string> {
   return invoke("test_connection", { config });
 }