feat(ai): add table mentions

Add table mentions to AI context, Ask/Agent mode selection, and related UI/i18n fixes.
2026-10-30 16:21:39 +00:00 · 2026-05-15 16:46:54 +08:00
parent 1563173173
commit 0099e1a262
13 changed files with 1150 additions and 49 deletions
@@ -11,6 +11,8 @@ import type {
 } from "@/types/database";
 import * as api from "@/lib/api";
 import { currentLocale } from "@/i18n";
+import { aiTableMentionKey, type AiTableMention } from "@/lib/aiTableMentions";
+import { isSchemaAware } from "@/lib/databaseCapabilities";

 export type AiAction = "generate" | "explain" | "optimize" | "fix" | "convert" | "sampleData";

@@ -169,6 +171,9 @@ export function buildSystemPrompt(action: AiAction, context: AiContext): string
      ? "当用户要求分析或查看某个表时，生成 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
      ? "对于 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.",
@@ -255,11 +260,12 @@ function formatSchema(context: AiContext): string {
 export async function buildAiContext(
  tab: QueryTab,
  connection: ConnectionConfig,
-  options: { maxTables?: number; maxColumnsPerTable?: number } = {},
+  options: { maxTables?: number; maxColumnsPerTable?: number; mentionedTables?: AiTableMention[] } = {},
 ): Promise<AiContext> {
  const maxTables = options.maxTables ?? 50;
  const maxColumnsPerTable = options.maxColumnsPerTable ?? 40;
  const tables: AiSchemaTable[] = [];
+  const tableKeys = new Set<string>();
  let truncated = false;

  if (tab.tableMeta) {
@@ -277,8 +283,20 @@ export async function buildAiContext(
      indexes,
      foreignKeys,
    });
+    tableKeys.add(aiTableMentionKey(tab.tableMeta.schema, tName));
    truncated = tab.tableMeta.columns.length > maxColumnsPerTable;
-  } else if (!["redis", "mongodb"].includes(connection.db_type)) {
+  }
+
+  for (const mention of options.mentionedTables ?? []) {
+    const key = aiTableMentionKey(mention.schema, mention.table);
+    if (tableKeys.has(key)) continue;
+    const entry = await loadMentionedTableContext(tab, connection, mention, maxColumnsPerTable).catch(() => undefined);
+    if (!entry) continue;
+    tableKeys.add(aiTableMentionKey(entry.schema, entry.name));
+    tables.push(entry);
+  }
+
+  if (!tab.tableMeta && !["redis", "mongodb"].includes(connection.db_type)) {
    try {
      const schemas = await loadCandidateSchemas(tab, connection);
      for (const schema of schemas) {
@@ -295,7 +313,7 @@ export async function buildAiContext(
                .listForeignKeys(tab.connectionId, tab.database, schema, table.name)
                .catch(() => [] as ForeignKeyInfo[]),
            ]).then(([columns, indexes, foreignKeys]) => ({
-              schema: schema === tab.database && connection.db_type !== "postgres" ? undefined : schema,
+              schema: schema === tab.database && !isSchemaAware(connection.db_type) ? undefined : schema,
              name: table.name,
              tableType: table.table_type,
              columns: columns.slice(0, maxColumnsPerTable),
@@ -309,6 +327,9 @@ export async function buildAiContext(
        for (const meta of metaResults) {
          if (meta._truncatedCols) truncated = true;
          const { _truncatedCols, ...entry } = meta;
+          const key = aiTableMentionKey(entry.schema, entry.name);
+          if (tableKeys.has(key)) continue;
+          tableKeys.add(key);
          tables.push(entry);
        }
        if (tables.length >= maxTables) break;
@@ -330,8 +351,49 @@ export async function buildAiContext(
  };
 }

+async function loadMentionedTableContext(
+  tab: QueryTab,
+  connection: ConnectionConfig,
+  mention: AiTableMention,
+  maxColumnsPerTable: number,
+): Promise<AiSchemaTable | undefined> {
+  const schema = await resolveMentionedTableSchema(tab, connection, mention);
+  const [columns, indexes, foreignKeys] = await Promise.all([
+    api.getColumns(tab.connectionId, tab.database, schema, mention.table),
+    api.listIndexes(tab.connectionId, tab.database, schema, mention.table).catch(() => [] as IndexInfo[]),
+    api.listForeignKeys(tab.connectionId, tab.database, schema, mention.table).catch(() => [] as ForeignKeyInfo[]),
+  ]);
+  return {
+    schema: schema === tab.database && !isSchemaAware(connection.db_type) ? undefined : schema,
+    name: mention.table,
+    tableType: "TABLE",
+    columns: columns.slice(0, maxColumnsPerTable),
+    indexes,
+    foreignKeys,
+  };
+}
+
+async function resolveMentionedTableSchema(
+  tab: QueryTab,
+  connection: ConnectionConfig,
+  mention: AiTableMention,
+): Promise<string> {
+  if (mention.schema) return mention.schema;
+  if (tab.tableMeta?.tableName.toLowerCase() === mention.table.toLowerCase() && tab.tableMeta.schema) {
+    return tab.tableMeta.schema;
+  }
+  if (isSchemaAware(connection.db_type)) {
+    const schemas = await loadCandidateSchemas(tab, connection);
+    for (const schema of schemas) {
+      const tables = await api.listTables(tab.connectionId, tab.database, schema, mention.table, 10).catch(() => []);
+      if (tables.some((table) => table.name.toLowerCase() === mention.table.toLowerCase())) return schema;
+    }
+  }
+  return tab.database || connection.database || "main";
+}
+
 async function loadCandidateSchemas(tab: QueryTab, connection: ConnectionConfig): Promise<string[]> {
-  if (connection.db_type === "postgres" || connection.db_type === "sqlserver") {
+  if (isSchemaAware(connection.db_type)) {
    const schemas = await api.listSchemas(tab.connectionId, tab.database);
    return prioritizeSchemas(schemas);
  }
@@ -0,0 +1,103 @@
+export interface AiTableMention {
+  raw: string;
+  schema?: string;
+  table: string;
+}
+
+const SIMPLE_IDENTIFIER_RE = /^[\p{L}_][\p{L}\p{N}_$]*$/u;
+
+function isMentionBoundary(char: string | undefined): boolean {
+  return !char || /\s|[([{,;:]/.test(char);
+}
+
+function isUnquotedMentionChar(char: string): boolean {
+  return /[\p{L}\p{N}_$.-]/u.test(char);
+}
+
+function readQuotedSegment(source: string, start: number): { value: string; end: number } | null {
+  const quote = source[start];
+  if (quote !== '"' && quote !== "`") return null;
+  let value = "";
+  for (let i = start + 1; i < source.length; i++) {
+    const char = source[i];
+    if (char === quote) {
+      if (source[i + 1] === quote) {
+        value += quote;
+        i++;
+        continue;
+      }
+      return { value, end: i + 1 };
+    }
+    value += char;
+  }
+  return null;
+}
+
+function readMentionToken(source: string, start: number): { raw: string; parts: string[]; end: number } | null {
+  let i = start;
+  const parts: string[] = [];
+  let raw = "";
+
+  while (i < source.length) {
+    const quoted = readQuotedSegment(source, i);
+    if (quoted) {
+      parts.push(quoted.value);
+      raw += source.slice(i, quoted.end);
+      i = quoted.end;
+    } else {
+      let value = "";
+      const valueStart = i;
+      while (i < source.length && isUnquotedMentionChar(source[i]) && source[i] !== ".") {
+        value += source[i];
+        i++;
+      }
+      if (!value) break;
+      parts.push(value);
+      raw += source.slice(valueStart, i);
+    }
+
+    if (source[i] !== ".") break;
+    raw += ".";
+    i++;
+  }
+
+  const usableParts = parts.filter(Boolean);
+  if (!usableParts.length) return null;
+  return { raw, parts: usableParts, end: i };
+}
+
+export function parseAiTableMentions(text: string): AiTableMention[] {
+  const mentions: AiTableMention[] = [];
+  const seen = new Set<string>();
+
+  for (let i = 0; i < text.length; i++) {
+    if (text[i] !== "@" || !isMentionBoundary(text[i - 1])) continue;
+    const token = readMentionToken(text, i + 1);
+    if (!token) continue;
+
+    const table = token.parts[token.parts.length - 1];
+    const schema = token.parts.length > 1 ? token.parts[token.parts.length - 2] : undefined;
+    const key = `${schema || ""}.${table}`.toLowerCase();
+    if (!seen.has(key)) {
+      seen.add(key);
+      mentions.push({ raw: `@${token.raw}`, schema, table });
+    }
+    i = token.end - 1;
+  }
+
+  return mentions;
+}
+
+export function formatAiTableMention(schema: string | undefined, table: string): string {
+  const parts = [schema, table].filter((part): part is string => !!part);
+  return `@${parts.map(formatMentionPart).join(".")}`;
+}
+
+function formatMentionPart(part: string): string {
+  if (SIMPLE_IDENTIFIER_RE.test(part)) return part;
+  return `"${part.replace(/"/g, '""')}"`;
+}
+
+export function aiTableMentionKey(schema: string | undefined, table: string): string {
+  return `${schema || ""}.${table}`.toLowerCase();
+}
@@ -29,13 +29,18 @@ export interface DataGridSaveStatementOptions {

 export interface DataGridSaveValidationOptions {
  databaseType?: DatabaseType;
+  tableMeta?: DataGridTableMeta;
  columns: string[];
+  rows?: GridCellValue[][];
  columnInfo?: DataGridColumnInfo[];
  dirtyRows: Array<[number, Array<[number, GridCellValue]>]>;
  newRows: GridCellValue[][];
 }

 export function validateDataGridSave(options: DataGridSaveValidationOptions): string | undefined {
+  const duplicateKeyError = validateInsertedPrimaryKeys(options);
+  if (duplicateKeyError) return duplicateKeyError;
+
  const notNullColumns = new Set(
    (options.columnInfo ?? [])
      .filter(
@@ -69,6 +74,32 @@ export function validateDataGridSave(options: DataGridSaveValidationOptions): st
  return undefined;
 }

+function validateInsertedPrimaryKeys(options: DataGridSaveValidationOptions): string | undefined {
+  const primaryKeys = options.tableMeta?.primaryKeys ?? [];
+  if (primaryKeys.length === 0 || options.newRows.length === 0) return undefined;
+
+  const primaryKeyIndexes = primaryKeys.map((primaryKey) => findColumnIndex(options.columns, primaryKey));
+  if (primaryKeyIndexes.some((index) => index === -1)) return undefined;
+
+  const existingKeys = new Set<string>();
+  for (const row of options.rows ?? []) {
+    const key = primaryKeyValueKey(primaryKeyIndexes, row);
+    if (key) existingKeys.add(key);
+  }
+
+  const newKeys = new Set<string>();
+  for (const row of options.newRows) {
+    const key = primaryKeyValueKey(primaryKeyIndexes, row);
+    if (!key) continue;
+    if (existingKeys.has(key) || newKeys.has(key)) {
+      return duplicatePrimaryKeyError(primaryKeys, primaryKeyIndexes, row, existingKeys.has(key));
+    }
+    newKeys.add(key);
+  }
+
+  return undefined;
+}
+
 export function buildDataGridSaveStatements(options: DataGridSaveStatementOptions): string[] {
  if (options.databaseType === "neo4j") return buildNeo4jDataGridSaveStatements(options);

@@ -252,6 +283,36 @@ function isNullWriteToNotNullColumn(
  return value === null && notNullColumns.has(normalizeColumnName(column));
 }

+function findColumnIndex(columns: string[], target: string): number {
+  const normalizedTarget = normalizeColumnName(target);
+  return columns.findIndex((column) => normalizeColumnName(column) === normalizedTarget);
+}
+
+function primaryKeyValueKey(primaryKeyIndexes: number[], row: GridCellValue[]): string | undefined {
+  const values = primaryKeyIndexes.map((index) => row[index]);
+  if (values.some((value) => value === null || value === undefined)) return undefined;
+  return JSON.stringify(values);
+}
+
+function duplicatePrimaryKeyError(
+  primaryKeys: string[],
+  primaryKeyIndexes: number[],
+  row: GridCellValue[],
+  matchesExistingRow: boolean,
+): string {
+  const keySummary = primaryKeys
+    .map((primaryKey, index) => `${primaryKey} = ${formatKeyValueForMessage(row[primaryKeyIndexes[index]])}`)
+    .join(", ");
+  const source = matchesExistingRow ? "the existing primary key" : "another new row's primary key";
+  return `New row duplicates ${source} (${keySummary}). Change the key before saving.`;
+}
+
+function formatKeyValueForMessage(value: GridCellValue): string {
+  if (value === null || value === undefined) return "NULL";
+  if (typeof value === "string") return `"${value.replace(/"/g, '\\"')}"`;
+  return String(value);
+}
+
 function normalizeColumnName(name: string): string {
  return name.toUpperCase();
 }