From 0e0edcd206bec15cb48f2c8acbd4c467f50560da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=8C=E4=B8=AB=E8=AE=B2=E6=A2=B5?= Date: Thu, 20 Aug 2026 15:15:19 +0800 Subject: [PATCH] feat(ai): add connection tree context --- apps/desktop/src/App.vue | 60 ++++++++++++++++++- .../src/components/editor/AiAssistant.vue | 22 ++++++- .../src/components/layout/AppSidebar.vue | 4 +- .../src/components/sidebar/ConnectionTree.vue | 2 + .../sidebar/SidebarTreeRuntimeHost.vue | 16 +++++ apps/desktop/src/i18n/locales/en.ts | 1 + apps/desktop/src/i18n/locales/es.ts | 1 + apps/desktop/src/i18n/locales/it.ts | 1 + apps/desktop/src/i18n/locales/ja.ts | 1 + apps/desktop/src/i18n/locales/ko.ts | 1 + apps/desktop/src/i18n/locales/pt-BR.ts | 1 + apps/desktop/src/i18n/locales/zh-CN.ts | 1 + apps/desktop/src/i18n/locales/zh-TW.ts | 1 + .../ai/addToAiDatabaseContext.spec.ts | 15 +++++ .../lib/database/databaseFeatureSupport.ts | 9 +++ 15 files changed, 132 insertions(+), 4 deletions(-) create mode 100644 apps/desktop/src/lib/__tests__/ai/addToAiDatabaseContext.spec.ts diff --git a/apps/desktop/src/App.vue b/apps/desktop/src/App.vue index 12eb36c540..c42910b85d 100644 --- a/apps/desktop/src/App.vue +++ b/apps/desktop/src/App.vue @@ -62,7 +62,7 @@ import { isTauriRuntime } from "@/lib/backend/tauriRuntime"; import { openQueryResultArchiveFile } from "@/lib/query/queryResultArchiveFile"; import { rememberExternalSqlFileTarget, resolveExternalSqlFileTarget, unassociatedExternalSqlFileTarget } from "@/lib/sql/externalSqlFileTarget"; import { externalSqlFileOpenErrorMessage, readBrowserSqlFile, sqlFileTitleFromPath } from "@/lib/sql/sqlFileOpen"; -import type { ConnectionConfig, DatabaseType, ObjectSourceKind, QueryTab } from "@/types/database"; +import type { ConnectionConfig, DatabaseType, ObjectSourceKind, QueryTab, TreeNode } from "@/types/database"; import { parseConnectionDeepLink, type ConnectionDeepLinkDraft } from "@/lib/connection/connectionDeepLink"; import { isBrowserReloadShortcut, @@ -109,6 +109,7 @@ import { savedSqlErrorMessage } from "@/lib/savedSql/savedSqlErrors"; import { savedSqlDefaultTargetForWrite } from "@/lib/savedSql/savedSqlExecutionTarget"; import { countActiveUpdateBlockingTasks } from "@/lib/app/appUpdateTaskGuard"; import { initSavedSqlEditorPositions } from "@/lib/app/savedSqlEditorPosition"; +import { hasTreeNodeDatabaseContext } from "@/lib/sidebar/treeNodeContext"; import { isSchemaAware, isSingleDatabase, usesTreeSchemaMode } from "@/lib/database/databaseFeatureSupport"; import { codeMirrorSqlDialect, connectionUsesDatabaseObjectTreeMode, effectiveDatabaseTypeForConnection } from "@/lib/database/jdbcDialect"; import { sqlFormatDialectForDbType } from "@/lib/sql/sqlFormatter"; @@ -120,7 +121,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { SearchableSelect } from "@/components/ui/searchable-select"; import type { HistoryEntry } from "@/lib/backend/tauri"; -import type { AiAction } from "@/lib/ai/ai"; +import { resolveDefaultAiSchema, type AiAction } from "@/lib/ai/ai"; import ExternalSqlFileChangeDialog from "@/components/editor/ExternalSqlFileChangeDialog.vue"; const AiAssistant = defineAsyncComponent(() => import("@/components/editor/AiAssistant.vue")); @@ -139,6 +140,8 @@ const QueryEditorObjectSourceDialog = defineAsyncComponent(() => import("@/compo type AiAssistantHandle = { triggerAction: (action: AiAction, instruction?: string) => void; setPrompt: (text: string) => void; + addTableMention: (target: { schema?: string; table: string }) => void; + clearContextReferences: () => void; }; const { t } = useI18n(); @@ -857,6 +860,58 @@ function sendSelectionToAi(sql: string) { invokeWhenAiReady((handle) => handle.setPrompt(sql)); } +let addToAiRequestId = 0; + +async function addToAi(node: TreeNode) { + if ((node.type !== "connection" && node.type !== "database" && node.type !== "table") || !node.connectionId) return; + const connection = connectionStore.getConfig(node.connectionId); + if (!connection) return; + const requestId = ++addToAiRequestId; + + try { + await connectionStore.ensureConnected(node.connectionId); + if (requestId !== addToAiRequestId) return; + connectionStore.activeConnectionId = node.connectionId; + + let target: { database: string; schema?: string; catalog?: string } | null = null; + if (node.type === "connection") { + const options = await getDatabaseOptions(node.connectionId); + if (requestId !== addToAiRequestId) return; + target = + connection.db_type === "dameng" + ? { + database: resolveDefaultDatabase(connection, []), + schema: resolveDefaultAiSchema(connection, options), + } + : { + database: resolveDefaultDatabase(connection, options), + schema: connection.default_schema, + }; + } else if (hasTreeNodeDatabaseContext(node)) { + target = { database: node.database, schema: node.schema, catalog: node.catalog }; + } + if (!target) return; + + const currentTab = activeTab.value; + const contextChanged = currentTab?.connectionId !== node.connectionId || currentTab?.database !== target.database || (currentTab?.schema || "") !== (target.schema || "") || (currentTab?.catalog || "") !== (target.catalog || ""); + + const contextTab = queryStore.tabs.find((tab) => tab.connectionId === node.connectionId && tab.database === target.database && (tab.schema || "") === (target.schema || "") && (tab.catalog || "") === (target.catalog || "")); + if (contextTab) { + queryStore.switchTab(contextTab.id); + } else { + queryStore.createTab(node.connectionId, target.database, undefined, "query", target.schema, undefined, target.catalog); + } + + openRightSidebarPanel("ai"); + invokeWhenAiReady((handle) => { + if (contextChanged) handle.clearContextReferences(); + if (node.type === "table") handle.addTableMention({ schema: node.schema, table: node.label }); + }); + } catch (e: any) { + toast(t("connection.connectFailed", { message: translateBackendError(t, e) }), 5000); + } +} + function openAiPanel() { openRightSidebarPanel("ai"); } @@ -2669,6 +2724,7 @@ onUnmounted(() => { @start-resize="startSidebarResize" @collapse="setSidebarOpen(false)" @open-settings="(initialTab) => openSettings(initialTab ?? 'appearance')" + @add-to-ai="addToAi" />
- +
diff --git a/apps/desktop/src/components/sidebar/ConnectionTree.vue b/apps/desktop/src/components/sidebar/ConnectionTree.vue index 5ec4c4243e..6a5dff66bb 100644 --- a/apps/desktop/src/components/sidebar/ConnectionTree.vue +++ b/apps/desktop/src/components/sidebar/ConnectionTree.vue @@ -97,6 +97,7 @@ const sidebarContextMenuRef = ref<{ close: () => void } | null>(null); const sidebarContextMenuItems = ref([]); const emit = defineEmits<{ "open-settings": [initialTab: string]; + "add-to-ai": [node: TreeNode]; }>(); const sidebarContextMenuTarget = ref(null); @@ -2159,6 +2160,7 @@ defineExpose({ focusSearch, createNewGroup, collapseAllTreeNodes }); @open-visible-schemas="openSidebarVisibleSchemas" @open-visible-nacos-namespaces="openSidebarVisibleNacosNamespaces" @open-table-name-filters="openSidebarTableNameFilters" + @add-to-ai="(node) => emit('add-to-ai', node)" @request-group-rename="startRenamingCreatedGroup" @request-saved-sql-rename="startRenamingSavedSqlNode" @open-danger-dialog="openSidebarDangerDialog" diff --git a/apps/desktop/src/components/sidebar/SidebarTreeRuntimeHost.vue b/apps/desktop/src/components/sidebar/SidebarTreeRuntimeHost.vue index 6eb463cb35..69ecda600e 100644 --- a/apps/desktop/src/components/sidebar/SidebarTreeRuntimeHost.vue +++ b/apps/desktop/src/components/sidebar/SidebarTreeRuntimeHost.vue @@ -57,6 +57,7 @@ import { X, Settings2, GitBranch, + Sparkles, } from "@lucide/vue"; import type { ContextMenuItem } from "@/components/ui/CustomContextMenu.vue"; import { CONNECTION_ATTEMPT_CANCELLED_MESSAGE, useConnectionStore } from "@/stores/connectionStore"; @@ -92,6 +93,7 @@ import { supportsDatabaseSearch, supportsConnectionDatabaseBrowser, supportsConnectionQueryActions, + supportsAiAssistantContext, supportsFieldLineage, supportsObjectBrowserTreeNode, supportsSchemaDiagram, @@ -374,6 +376,7 @@ const emit = defineEmits<{ "open-visible-schemas": [node: TreeNode]; "open-visible-nacos-namespaces": [node: TreeNode]; "open-table-name-filters": [node: TreeNode]; + "add-to-ai": [node: TreeNode]; "open-danger-dialog": [request: SidebarDangerDialogRequest]; "open-dialog-controller": [controller: Record | null]; "open-install-extension": [node: TreeNode]; @@ -4581,8 +4584,12 @@ function buildConnectionSidebarMenu(context: SidebarMenuFactoryContext): boolean items.push({ label: t("contextMenu.copyName"), action: copyName, icon: Copy, shortcut: shortcutCopyName.value }); items.push({ label: "", separator: true }); const supportsQueryActions = supportsConnectionQueryActions(currentDatabaseType()); + const supportsAiContext = supportsAiAssistantContext(currentDatabaseType()); if (supportsQueryActions) { items.push({ label: t("contextMenu.newQuery"), action: newQuery, icon: TerminalSquare }); + if (supportsAiContext) { + items.push({ label: t("contextMenu.addToAi"), action: () => emit("add-to-ai", node), icon: Sparkles }); + } } const connectionWorkspace = node.connectionId ? driverProfileDatabaseWorkspace(connectionStore.getConfig(node.connectionId)?.driver_profile) : undefined; if (connectionWorkspace?.entryScopes.includes("connection")) { @@ -4777,6 +4784,9 @@ function buildDatabaseSidebarMenu(context: SidebarMenuFactoryContext): boolean { } if (supportsConnectionQueryActions(currentDatabaseType())) { items.push({ label: t("contextMenu.newQuery"), action: newQuery, icon: TerminalSquare }); + if (node.type === "database" && supportsAiAssistantContext(currentDatabaseType())) { + items.push({ label: t("contextMenu.addToAi"), action: () => emit("add-to-ai", node), icon: Sparkles }); + } const sqlHistoryMenu = savedSqlHistorySubmenu(); if (sqlHistoryMenu) items.push(sqlHistoryMenu); } @@ -5072,6 +5082,9 @@ function buildObjectSidebarMenu(context: SidebarMenuFactoryContext): boolean { items.push({ label: t("contextMenu.viewData"), action: openDataImmediately, icon: TableProperties }); items.push({ label: t("contextMenu.openInNewDataTab"), action: openDataInNewTabImmediately, icon: CopyPlus }); items.push({ label: t("contextMenu.newQuery"), action: newQuery, icon: TerminalSquare }); + if (supportsAiAssistantContext(currentDatabaseType())) { + items.push({ label: t("contextMenu.addToAi"), action: () => emit("add-to-ai", node), icon: Sparkles }); + } items.push({ label: "", separator: true }); items.push(exportDataSubmenu(false)); items.push({ label: t("contextMenu.refreshChildren"), action: refresh, icon: RefreshCw, shortcut: shortcutRefresh }); @@ -5080,6 +5093,9 @@ function buildObjectSidebarMenu(context: SidebarMenuFactoryContext): boolean { const destructiveActions: ContextMenuItem[] = []; items.push(copyNameMenuItem()); items.push({ label: t("contextMenu.newQuery"), action: newQuery, icon: TerminalSquare }); + if (node.type === "table" && supportsAiAssistantContext(currentDatabaseType())) { + items.push({ label: t("contextMenu.addToAi"), action: () => emit("add-to-ai", node), icon: Sparkles }); + } items.push({ label: "", separator: true }); items.push({ label: t("contextMenu.viewData"), action: openDataImmediately, icon: TableProperties }); items.push({ diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 47b630a15b..b7aeced89c 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -2698,6 +2698,7 @@ export default { duplicateConnection: "Duplicate Connection", duplicateSelectedConnections: "Duplicate Selected {count} Connections", newQuery: "New Query", + addToAi: "Add to AI", instanceInfo: "Instance Info", newSql: "New SQL", generateSql: "Generate SQL", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index 763d7ac504..a5a21c350e 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -2631,6 +2631,7 @@ export default withEnglishFallback({ duplicateSelectedConnections: "Duplicar {count} conexiones seleccionadas", instanceInfo: "Información de la instancia", newQuery: "Nueva consulta", + addToAi: "Agregar a la IA", newSql: "Nuevo SQL", generateSql: "Generar SQL", newInsert: "Nueva Inserción", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index f8613122e7..88b463a7cf 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -2628,6 +2628,7 @@ export default withEnglishFallback({ duplicateConnection: "Duplica Connessione", duplicateSelectedConnections: "Duplica {count} connessioni selezionate", newQuery: "Nuova Query", + addToAi: "Aggiungi all'IA", instanceInfo: "Informazioni Istanza", newSql: "Nuovo SQL", generateSql: "Genera SQL", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index 0b99734292..c3fc360fda 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -2654,6 +2654,7 @@ export default withEnglishFallback({ duplicateConnection: "接続を複製", duplicateSelectedConnections: "選択した{count}件の接続を複製", newQuery: "新しいクエリ", + addToAi: "AI に追加", newSql: "新しいSQL", generateSql: "SQLを生成", newInsert: "新しいINSERT", diff --git a/apps/desktop/src/i18n/locales/ko.ts b/apps/desktop/src/i18n/locales/ko.ts index b7278b1c3d..8f03d91442 100644 --- a/apps/desktop/src/i18n/locales/ko.ts +++ b/apps/desktop/src/i18n/locales/ko.ts @@ -2542,6 +2542,7 @@ export default withEnglishFallback({ duplicateConnection: "연결 복제", duplicateSelectedConnections: "선택한 연결 {count}개 복제", newQuery: "새 쿼리", + addToAi: "AI에 추가", instanceInfo: "인스턴스 정보", newSql: "새 SQL", generateSql: "SQL 생성", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index 84fe6b9177..2f98d77250 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -2631,6 +2631,7 @@ export default withEnglishFallback({ duplicateSelectedConnections: "Duplicar {count} conexões selecionadas", instanceInfo: "Informações da Instância", newQuery: "Nova Consulta", + addToAi: "Adicionar à IA", newSql: "Novo SQL", generateSql: "Gerar SQL", newInsert: "Novo Insert", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index 5988a992b3..e4e06fb2c2 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -2622,6 +2622,7 @@ export default withEnglishFallback({ damengJobAdmin: "达梦代理作业", openDamengJobAdmin: "打开达梦代理作业", newQuery: "新建查询", + addToAi: "添加到 AI", instanceInfo: "实例信息", newInsert: "新建新增", newSql: "新建 SQL", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index 17f85e2920..b03d24716d 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -2629,6 +2629,7 @@ export default withEnglishFallback({ duplicateConnection: "複製連線", duplicateSelectedConnections: "複製選取的 {count} 個連線", newQuery: "建立查詢", + addToAi: "加入 AI", instanceInfo: "實例資訊", newSql: "新增 SQL", generateSql: "產生 SQL", diff --git a/apps/desktop/src/lib/__tests__/ai/addToAiDatabaseContext.spec.ts b/apps/desktop/src/lib/__tests__/ai/addToAiDatabaseContext.spec.ts new file mode 100644 index 0000000000..e736c63510 --- /dev/null +++ b/apps/desktop/src/lib/__tests__/ai/addToAiDatabaseContext.spec.ts @@ -0,0 +1,15 @@ +import { readFileSync } from "node:fs"; +import { describe, expect, it } from "vitest"; + +const appSource = readFileSync(new URL("../../../App.vue", import.meta.url), "utf8"); + +describe("add to AI database context", () => { + it("accepts database tree nodes with an empty database identifier", () => { + const addToAiStart = appSource.indexOf("async function addToAi"); + const addToAiEnd = appSource.indexOf("function openAiPanel", addToAiStart); + const addToAiSource = appSource.slice(addToAiStart, addToAiEnd); + + expect(addToAiSource).toContain("else if (hasTreeNodeDatabaseContext(node))"); + expect(addToAiSource).not.toContain("else if (node.database)"); + }); +}); diff --git a/apps/desktop/src/lib/database/databaseFeatureSupport.ts b/apps/desktop/src/lib/database/databaseFeatureSupport.ts index 860bf49f3a..0eb43f6ccb 100644 --- a/apps/desktop/src/lib/database/databaseFeatureSupport.ts +++ b/apps/desktop/src/lib/database/databaseFeatureSupport.ts @@ -92,6 +92,15 @@ export function supportsQueryExecution(dbType?: DatabaseType): boolean { return supportsDatabaseFeature(dbType, "queryExecution"); } +/** + * The AI assistant currently builds its context from database/table metadata. + * Connection-only query targets (for example etcd and ZooKeeper) do not expose + * that hierarchy, so they must not be offered by sidebar "Add to AI" actions. + */ +export function supportsAiAssistantContext(dbType?: DatabaseType): boolean { + return supportsQueryExecution(dbType) && !usesConnectionOnlyQueryTarget(dbType); +} + export function supportsConnectionScopedQueryExecution(dbType?: DatabaseType): boolean { return supportsRegisteredConnectionScopedQueryExecution(dbType); }