From d90847bbba14bd1b03581823376865f744747e24 Mon Sep 17 00:00:00 2001 From: zipg Date: Tue, 7 Jul 2026 15:39:10 +0800 Subject: [PATCH] feat(ai): support SQL library file references --- .../src/components/editor/AiAssistant.vue | 748 ++++++++++++++---- apps/desktop/src/i18n/locales/en.ts | 5 +- apps/desktop/src/i18n/locales/es.ts | 5 +- apps/desktop/src/i18n/locales/it.ts | 5 +- apps/desktop/src/i18n/locales/ja.ts | 5 +- apps/desktop/src/i18n/locales/pt-BR.ts | 5 +- apps/desktop/src/i18n/locales/zh-CN.ts | 5 +- apps/desktop/src/i18n/locales/zh-TW.ts | 5 +- apps/desktop/src/lib/__tests__/ai/ai.spec.ts | 1 + apps/desktop/src/lib/ai/ai.ts | 29 +- apps/desktop/src/lib/backend/tauri.ts | 1 + crates/dbx-core/src/ai.rs | 2 + packages/app-tests/aiPrompt.test.ts | 25 + packages/app-tests/aiPromptEval.test.ts | 2 + packages/app-tests/aiSchemaComments.test.ts | 2 + 15 files changed, 667 insertions(+), 178 deletions(-) diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue index 726fe2b581..15206d9e6d 100644 --- a/apps/desktop/src/components/editor/AiAssistant.vue +++ b/apps/desktop/src/components/editor/AiAssistant.vue @@ -4,6 +4,7 @@ import { uuid } from "@/lib/common/utils"; import { useI18n } from "vue-i18n"; import { translateBackendError } from "@/i18n/backend-errors"; import { + ArrowDown, ArrowUp, ArrowRightLeft, AlertTriangle, @@ -13,6 +14,7 @@ import { CircleSlash, Copy, Database, + FileCode, GitBranch, HelpCircle, History, @@ -44,11 +46,13 @@ import { useTheme } from "@/composables/useTheme"; import { useSettingsStore, AI_PROVIDER_PRESETS, type AiProvider } from "@/stores/settingsStore"; import AiProviderLogo from "@/components/icons/AiProviderLogo.vue"; import { useConnectionStore } from "@/stores/connectionStore"; +import { useSavedSqlStore } from "@/stores/savedSqlStore"; import { connectionIconType } from "@/lib/connection/connectionPresentation"; import DatabaseIcon from "@/components/icons/DatabaseIcon.vue"; import { useQueryStore } from "@/stores/queryStore"; import { useToast } from "@/composables/useToast"; -import { buildAiContext, runAgentStream, isVectorDbType, defaultActionForMode, isValidActionForMode, type AiAction, type AiAssistantMode } from "@/lib/ai/ai"; +import { useNavigationTargets } from "@/composables/useNavigationTargets"; +import { buildAiContext, runAgentStream, isVectorDbType, defaultActionForMode, isValidActionForMode, type AiAction, type AiAssistantMode, type AiSqlFileContext } from "@/lib/ai/ai"; import { formatAiModelOption } from "@/lib/ai/aiModelPresentation"; import type { AgentEvent } from "@/lib/backend/tauri"; import { buildAiAgentPlan } from "@/lib/ai/aiAgentPlan"; @@ -58,7 +62,7 @@ import { createAiMessageRenderer } from "@/lib/ai/aiMessageRender"; import { formatAiInlineMarkdown, handleAiMarkdownLinkClick } from "@/lib/ai/aiMarkdown"; import { aiCancelStream, aiListModels, saveAiConversation, loadAiConversations, deleteAiConversation, listSchemas, listTables, type AiConversation, type AiModelInfo } from "@/lib/backend/api"; import type { AiMessage } from "@/lib/backend/api"; -import type { ConnectionConfig, QueryTab, TableInfo } from "@/types/database"; +import type { ConnectionConfig, QueryTab, SavedSqlFile, TableInfo } from "@/types/database"; import { useDatabaseOptions } from "@/composables/useDatabaseOptions"; import { decodeSelectableDatabaseValue, encodeSelectableDatabaseValue, formatDatabaseLabel, resolveDefaultDatabase } from "@/lib/database/defaultDatabase"; import { isSchemaAware } from "@/lib/database/databaseCapabilities"; @@ -73,13 +77,36 @@ import { visibleToActualIndex } from "@/lib/ai/aiMessageEdit"; const { t } = useI18n(); const settings = useSettingsStore(); const connectionStore = useConnectionStore(); +const savedSqlStore = useSavedSqlStore(); const queryStore = useQueryStore(); +const { openTableTarget } = useNavigationTargets({ + showFieldLineageDialog: ref(false), + showDatabaseSearchDialog: ref(false), +}); const { toast } = useToast(); const { isDark } = useTheme(); +type AiMessageMention = + | { + kind: "table"; + raw: string; + connectionId: string; + database: string; + schema?: string; + table: string; + } + | { + kind: "sqlFile"; + raw: string; + connectionId: string; + id: string; + name: string; + }; + interface ChatMessage { role: "user" | "assistant"; content: string; + mentions?: AiMessageMention[]; reasoning?: string; isThinking?: boolean; agentSteps?: AiAgentStepItem[]; @@ -111,6 +138,9 @@ const conversationId = ref(""); const conversations = ref([]); const showConversationList = ref(false); const promptTextareaRef = ref(null); +const shouldAutoScroll = ref(true); +const userPausedAutoScroll = ref(false); +const showScrollToBottom = ref(false); const promptCompositionActive = ref(false); const shikiCodeHighlighter = ref(); const agentTokens = ref<{ input: number; output: number } | null>(null); @@ -120,12 +150,21 @@ const draftBeforeHistory = ref(""); const editingMessageIndex = ref(null); const editingContent = ref(""); +const editingMentions = ref([]); const editCompositionActive = ref(false); +const MESSAGE_SCROLL_RESUME_THRESHOLD_PX = 16; +const MESSAGE_SCROLL_BUTTON_SHOW_THRESHOLD_PX = 120; +const MESSAGE_SCROLL_BUTTON_HIDE_THRESHOLD_PX = 48; +let messageScrollViewport: HTMLElement | null = null; +let messageTouchStartY: number | null = null; +let lastMessageScrollTop = 0; function startEditMessage(visibleIndex: number) { if (isGenerating.value) return; editingMessageIndex.value = visibleIndex; - editingContent.value = visibleMessages.value[visibleIndex].content; + const msg = visibleMessages.value[visibleIndex]; + editingContent.value = msg.content; + editingMentions.value = promptMentionChipsFromMessage(msg); nextTick(() => { const el = document.querySelector("[data-edit-textarea]"); if (el) { @@ -138,11 +177,12 @@ function startEditMessage(visibleIndex: number) { function cancelEdit() { editingMessageIndex.value = null; editingContent.value = ""; + editingMentions.value = []; } function submitEdit(visibleIndex: number) { const content = editingContent.value.trim(); - if (!content) return; + if (!content && !editingMentions.value.length) return; const actualIndex = visibleToActualIndex(messages.value, visibleIndex); if (actualIndex < 0) return; if (!props.connection || !props.tab) return; @@ -153,7 +193,9 @@ function submitEdit(visibleIndex: number) { messages.value = messages.value.slice(0, actualIndex); editingMessageIndex.value = null; editingContent.value = ""; - selectedMentions.value = []; + selectedMentions.value = editingMentions.value.filter((mention): mention is AiTableMention & { kind: "table" } => mention.kind === "table").map(({ raw, schema, table }) => ({ raw, schema, table })); + selectedSqlFileMentions.value = editingMentions.value.filter((mention): mention is AiSqlFileMention => mention.kind === "sqlFile"); + editingMentions.value = []; prompt.value = content; send(); } @@ -254,12 +296,31 @@ let resizeStartY = 0; let resizeStartHeight = 0; let promptPanelResizeObserver: ResizeObserver | undefined; -interface AiMentionCandidate { +interface AiTableMentionCandidate { + kind: "table"; schema?: string; name: string; tableType: string; } +interface AiSqlFileMentionCandidate { + kind: "sqlFile"; + id: string; + name: string; + folderPath?: string; +} + +type AiMentionCandidate = AiTableMentionCandidate | AiSqlFileMentionCandidate; + +interface AiSqlFileMention { + kind: "sqlFile"; + raw: string; + id: string; + name: string; +} + +type AiPromptMentionChip = (AiTableMention & { kind: "table" }) | AiSqlFileMention; + const mentionOpen = ref(false); const mentionLoading = ref(false); const mentionError = ref(""); @@ -269,9 +330,13 @@ const mentionCandidates = ref([]); const mentionCache = ref>({}); const mentionListRef = ref(null); const selectedMentions = ref([]); +const selectedSqlFileMentions = ref([]); let mentionTimer: ReturnType | undefined; let mentionRequestId = 0; +const AI_SQL_FILE_MENTION_CANDIDATE_LIMIT = 50; +const AI_SQL_FILE_CONTEXT_MAX_CHARS = 12_000; + interface AiActionButton { action: AiAction; icon: Component; @@ -348,22 +413,35 @@ function messagesForAgentHistory(historyMessages: ChatMessage[]): AiMessage[] { } } if (latestSummaryIndex < 0) { - return historyMessages.map((m) => ({ role: m.role, content: m.content })); + return historyMessages.map((m) => ({ role: m.role, content: messageContentForModel(m) })); } const compactedHistory = historyMessages.slice(latestSummaryIndex); const firstMsg = historyMessages[0]; if (firstMsg && firstMsg.role === "user" && firstMsg.kind !== "contextSummary") { - return [{ role: "user" as const, content: firstMsg.content }, ...compactedHistory.map((m) => ({ role: m.role, content: m.content }))]; + return [{ role: "user" as const, content: messageContentForModel(firstMsg) }, ...compactedHistory.map((m) => ({ role: m.role, content: messageContentForModel(m) }))]; } - return compactedHistory.map((m) => ({ role: m.role, content: m.content })); + return compactedHistory.map((m) => ({ role: m.role, content: messageContentForModel(m) })); } const chatTitle = computed(() => { const first = messages.value.find((m) => m.role === "user" && m.kind !== "contextSummary"); - return first ? first.content.slice(0, 30) : t("ai.newChat"); + return first ? messageTitle(first).slice(0, 30) : t("ai.newChat"); }); -const promptMentionChips = computed(() => selectedMentions.value); +const promptMentionChips = computed(() => [...selectedMentions.value.map((mention) => ({ ...mention, kind: "table" as const })), ...selectedSqlFileMentions.value]); + +function messageMentionLabels(message: ChatMessage): string[] { + return promptMentionChipsFromMessage(message).map((mention) => mention.raw); +} + +function messageContentForModel(message: ChatMessage): string { + if (message.kind === "contextSummary") return message.content; + return [...messageMentionLabels(message), message.content].filter(Boolean).join(" "); +} + +function messageTitle(message: ChatMessage): string { + return [promptMentionChipsFromMessage(message).map(mentionDisplayName).join(" "), message.content].filter(Boolean).join(" ") || t("ai.newChat"); +} const isWaitingForFirstDelta = computed(() => { const last = messages.value[messages.value.length - 1]; @@ -511,7 +589,7 @@ function appendAssistantReasoning(assistantIdx: number, delta: string) { scrollToBottom(); } -const reasoningExpanded = ref(false); +const expandedReasoning = ref>(new Set()); const expandedSteps = ref>(new Set()); function toggleStep(key: string) { @@ -626,23 +704,169 @@ function agentEventToStep(event: AgentEvent, index: number): AiAgentStepItem | u }; } -function toggleReasoning() { - reasoningExpanded.value = !reasoningExpanded.value; +function toggleReasoning(index: number) { + const next = new Set(expandedReasoning.value); + if (next.has(index)) { + next.delete(index); + } else { + next.add(index); + } + expandedReasoning.value = next; } -function scrollToBottom() { +function getMessageScrollViewport(): HTMLElement | null { + const root = scrollRef.value?.$el as HTMLElement | undefined; + return root?.querySelector('[data-slot="scroll-area-viewport"]') as HTMLElement | null; +} + +function messageBottomDistance(el: HTMLElement) { + return Math.max(0, el.scrollHeight - el.scrollTop - el.clientHeight); +} + +function isAtMessageBottom(el: HTMLElement) { + return messageBottomDistance(el) <= MESSAGE_SCROLL_RESUME_THRESHOLD_PX; +} + +function messageCanScroll(el: HTMLElement) { + return el.scrollHeight > el.clientHeight + MESSAGE_SCROLL_RESUME_THRESHOLD_PX; +} + +function shouldShowMessageScrollButton(el: HTMLElement) { + if (!messageCanScroll(el)) return false; + const distance = messageBottomDistance(el); + return distance > (showScrollToBottom.value ? MESSAGE_SCROLL_BUTTON_HIDE_THRESHOLD_PX : MESSAGE_SCROLL_BUTTON_SHOW_THRESHOLD_PX); +} + +function updateMessageScrollButtonVisibility() { + const el = getMessageScrollViewport(); + showScrollToBottom.value = !!el && shouldShowMessageScrollButton(el); +} + +function pauseMessageAutoScroll() { + userPausedAutoScroll.value = true; + shouldAutoScroll.value = false; + updateMessageScrollButtonVisibility(); +} + +function updateMessageScrollState() { + const el = getMessageScrollViewport(); + if (!el) { + showScrollToBottom.value = false; + return; + } + if (isAtMessageBottom(el)) { + userPausedAutoScroll.value = false; + shouldAutoScroll.value = true; + showScrollToBottom.value = false; + return; + } + if (userPausedAutoScroll.value) { + shouldAutoScroll.value = false; + showScrollToBottom.value = shouldShowMessageScrollButton(el); + return; + } + shouldAutoScroll.value = false; + showScrollToBottom.value = shouldShowMessageScrollButton(el); +} + +function handleMessageScroll() { + const el = getMessageScrollViewport(); + if (!el) return; + if (el.scrollTop < lastMessageScrollTop - 2) { + userPausedAutoScroll.value = true; + } + lastMessageScrollTop = el.scrollTop; + updateMessageScrollState(); +} + +function handleMessageWheel(event: WheelEvent) { + if (event.deltaY < 0) pauseMessageAutoScroll(); +} + +function handleMessageTouchStart(event: TouchEvent) { + messageTouchStartY = event.touches[0]?.clientY ?? null; +} + +function handleMessageTouchMove(event: TouchEvent) { + if (messageTouchStartY == null) return; + const currentY = event.touches[0]?.clientY ?? messageTouchStartY; + if (currentY - messageTouchStartY > 4) pauseMessageAutoScroll(); +} + +function handleMessageKeydown(event: KeyboardEvent) { + if (["ArrowUp", "PageUp", "Home"].includes(event.key)) pauseMessageAutoScroll(); +} + +function detachMessageScrollListener() { + if (!messageScrollViewport) return; + messageScrollViewport.removeEventListener("scroll", handleMessageScroll); + messageScrollViewport.removeEventListener("wheel", handleMessageWheel); + messageScrollViewport.removeEventListener("touchstart", handleMessageTouchStart); + messageScrollViewport.removeEventListener("touchmove", handleMessageTouchMove); + messageScrollViewport.removeEventListener("keydown", handleMessageKeydown); + messageScrollViewport = null; +} + +function attachMessageScrollListener() { nextTick(() => { - const root = scrollRef.value?.$el as HTMLElement | undefined; - const el = root?.querySelector('[data-slot="scroll-area-viewport"]') as HTMLElement | null; + const el = getMessageScrollViewport(); + if (el === messageScrollViewport) return; + detachMessageScrollListener(); + messageScrollViewport = el; + if (!el) return; + el.addEventListener("scroll", handleMessageScroll, { passive: true }); + el.addEventListener("wheel", handleMessageWheel, { passive: true }); + el.addEventListener("touchstart", handleMessageTouchStart, { passive: true }); + el.addEventListener("touchmove", handleMessageTouchMove, { passive: true }); + el.addEventListener("keydown", handleMessageKeydown); + lastMessageScrollTop = el.scrollTop; + updateMessageScrollState(); + }); +} + +function scrollToBottom(options: { force?: boolean } = {}) { + if (options.force) { + userPausedAutoScroll.value = false; + shouldAutoScroll.value = true; + } + if (!options.force && (userPausedAutoScroll.value || !shouldAutoScroll.value)) { + updateMessageScrollButtonVisibility(); + return; + } + nextTick(() => { + const el = getMessageScrollViewport(); if (!el) return; requestAnimationFrame(() => { + if (!options.force && (userPausedAutoScroll.value || !shouldAutoScroll.value)) { + updateMessageScrollButtonVisibility(); + return; + } el.scrollTop = el.scrollHeight; + lastMessageScrollTop = el.scrollTop; + userPausedAutoScroll.value = false; + shouldAutoScroll.value = true; + showScrollToBottom.value = false; }); }); } +watch( + () => messages.value.length, + (length) => { + if (length) { + attachMessageScrollListener(); + return; + } + detachMessageScrollListener(); + userPausedAutoScroll.value = false; + shouldAutoScroll.value = true; + showScrollToBottom.value = false; + }, + { flush: "post" }, +); + function mentionCacheKey(connectionId: string, database: string, query: string) { - return `${connectionId}:${database}:${query.toLowerCase()}`; + return `${connectionId}:${database}:${savedSqlStore.version}:${query.toLowerCase()}`; } function mentionSchemaOrder(schemas: string[]): string[] { @@ -688,10 +912,12 @@ async function loadMentionCandidates(query: string) { mentionLoading.value = true; mentionError.value = ""; const { schemaPrefix, tableFilter } = normalizeMentionQuery(query); + let sqlFileCandidates: AiSqlFileMentionCandidate[] = []; try { + sqlFileCandidates = await loadSqlFileMentionCandidates(query); await connectionStore.ensureConnected(props.tab.connectionId); - let candidates: AiMentionCandidate[] = []; + let tableCandidates: AiMentionCandidate[] = []; if (isSchemaAware(props.connection.db_type)) { const schemas = mentionSchemaOrder(await listSchemas(props.tab.connectionId, props.tab.database)); const filteredSchemas = schemaPrefix ? schemas.filter((schema) => schema.toLowerCase().includes(schemaPrefix.toLowerCase())) : schemas; @@ -705,11 +931,11 @@ async function loadMentionCandidates(query: string) { ); }), ); - candidates = filterAiTableMentionCandidates(results.flat(), "", AI_TABLE_MENTION_CANDIDATE_LIMIT); + tableCandidates = filterAiTableMentionCandidates(results.flat(), "", AI_TABLE_MENTION_CANDIDATE_LIMIT); } else { const schema = props.tab.database || props.connection.database || "main"; const tables = await listTables(props.tab.connectionId, props.tab.database, schema, tableFilter || undefined, AI_TABLE_MENTION_CANDIDATE_LIMIT); - candidates = filterAiTableMentionCandidates( + tableCandidates = filterAiTableMentionCandidates( tables.map((table) => mentionCandidateFromTable(table)), tableFilter, AI_TABLE_MENTION_CANDIDATE_LIMIT, @@ -717,11 +943,18 @@ async function loadMentionCandidates(query: string) { } if (requestId !== mentionRequestId) return; - mentionCache.value[key] = candidates.slice(0, AI_TABLE_MENTION_CANDIDATE_LIMIT); + mentionCache.value[key] = [...tableCandidates, ...sqlFileCandidates]; mentionCandidates.value = mentionCache.value[key]; setMentionSelectedIndex(0); } catch (e: unknown) { if (requestId !== mentionRequestId) return; + if (sqlFileCandidates.length) { + mentionCache.value[key] = sqlFileCandidates; + mentionCandidates.value = sqlFileCandidates; + mentionError.value = ""; + setMentionSelectedIndex(0); + return; + } const message = e instanceof Error ? e.message : String(e); mentionError.value = translateBackendError(t, message); mentionCandidates.value = []; @@ -730,26 +963,142 @@ async function loadMentionCandidates(query: string) { } } -function mentionCandidateFromTable(table: TableInfo, schema?: string): AiMentionCandidate { - return { schema, name: table.name, tableType: table.table_type }; +async function loadSqlFileMentionCandidates(query: string): Promise { + const connectionId = props.tab?.connectionId; + if (!connectionId) return []; + await savedSqlStore.initFromStorage(); + const normalizedQuery = normalizeSqlFileMentionQuery(query); + return savedSqlStore.allFiles + .filter((file) => file.connectionId === connectionId) + .map((file) => ({ file, folderPath: savedSqlFolderPath(file) })) + .filter(({ file, folderPath }) => sqlFileMatchesQuery(file, folderPath, normalizedQuery)) + .slice(0, AI_SQL_FILE_MENTION_CANDIDATE_LIMIT) + .map(({ file, folderPath }) => ({ + kind: "sqlFile", + id: file.id, + name: file.name, + folderPath, + })); } -function mentionDisplayName(mention: AiTableMention) { +function normalizeSqlFileMentionQuery(query: string) { + return query.replace(/^["`{]+|["`}]+$/g, "").toLowerCase(); +} + +function sqlFileMatchesQuery(file: SavedSqlFile, folderPath: string | undefined, query: string) { + if (!query) return true; + return [file.name, folderPath || ""].some((value) => value.toLowerCase().includes(query)); +} + +function savedSqlFolderPath(file: SavedSqlFile): string | undefined { + if (!file.folderId) return undefined; + const foldersById = new Map(savedSqlStore.allFolders.map((folder) => [folder.id, folder])); + const names: string[] = []; + let current = foldersById.get(file.folderId); + while (current) { + names.unshift(current.name); + current = current.parentFolderId ? foldersById.get(current.parentFolderId) : undefined; + } + return names.length ? names.join(" / ") : undefined; +} + +function mentionCandidateFromTable(table: TableInfo, schema?: string): AiTableMentionCandidate { + return { kind: "table", schema, name: table.name, tableType: table.table_type }; +} + +function mentionCandidateName(candidate: AiMentionCandidate) { + if (candidate.kind === "sqlFile") return candidate.name; + return [candidate.schema, candidate.name].filter(Boolean).join("."); +} + +function mentionDisplayName(mention: AiPromptMentionChip) { + if (mention.kind === "sqlFile") return mention.name; return [mention.schema, mention.table].filter(Boolean).join("."); } -function removeMentionChip(mention: AiTableMention) { - selectedMentions.value = selectedMentions.value.filter((item) => item.raw !== mention.raw); +function promptMentionChipsFromMessage(message: ChatMessage): AiPromptMentionChip[] { + return (message.mentions || []).map((mention) => { + if (mention.kind === "sqlFile") return { kind: "sqlFile", raw: mention.raw, id: mention.id, name: mention.name }; + return { kind: "table", raw: mention.raw, schema: mention.schema, table: mention.table }; + }); +} + +function removeMentionChip(mention: AiPromptMentionChip) { + if (mention.kind === "sqlFile") { + selectedSqlFileMentions.value = selectedSqlFileMentions.value.filter((item) => item.id !== mention.id); + } else { + selectedMentions.value = selectedMentions.value.filter((item) => item.raw !== mention.raw); + } nextTick(() => promptTextareaRef.value?.focus()); } +function removeEditingMentionChip(index: number) { + editingMentions.value = editingMentions.value.filter((_, itemIndex) => itemIndex !== index); + nextTick(() => { + const el = document.querySelector("[data-edit-textarea]"); + el?.focus(); + }); +} + function addSelectedMention(candidate: AiMentionCandidate) { + if (candidate.kind === "sqlFile") { + const raw = `@{${candidate.name}}`; + if (selectedSqlFileMentions.value.some((mention) => mention.id === candidate.id)) return; + selectedSqlFileMentions.value.push({ kind: "sqlFile", raw, id: candidate.id, name: candidate.name }); + return; + } const raw = formatAiTableMention(candidate.schema, candidate.name); const key = `${candidate.schema || ""}.${candidate.name}`.toLowerCase(); if (selectedMentions.value.some((mention) => `${mention.schema || ""}.${mention.table}`.toLowerCase() === key)) return; selectedMentions.value.push({ raw, schema: candidate.schema, table: candidate.name }); } +function formatMentionCandidateType(candidate: AiMentionCandidate) { + if (candidate.kind === "sqlFile") return candidate.folderPath || "SQL"; + return formatMentionTableType(candidate.tableType); +} + +function selectedMessageMentions(tableMentions: AiTableMention[], sqlFileMentions: AiSqlFileMention[]): AiMessageMention[] { + const connectionId = props.tab?.connectionId || props.connection?.id || ""; + const database = props.tab?.database || props.connection?.database || ""; + return [ + ...tableMentions.map((mention) => ({ + kind: "table" as const, + raw: mention.raw, + connectionId, + database, + schema: mention.schema, + table: mention.table, + })), + ...sqlFileMentions.map((mention) => ({ + kind: "sqlFile" as const, + raw: mention.raw, + connectionId, + id: mention.id, + name: mention.name, + })), + ]; +} + +async function openMessageMention(mention: AiMessageMention) { + try { + if (mention.kind === "sqlFile") { + const file = await savedSqlStore.ensureFileContent(mention.id); + if (file) queryStore.openSavedSql(file); + return; + } + await openTableTarget({ + connectionId: mention.connectionId || props.tab?.connectionId || props.connection?.id || "", + database: mention.database || props.tab?.database || props.connection?.database || "", + schema: mention.schema, + tableName: mention.table, + }); + } catch (e: unknown) { + const message = e instanceof Error ? e.message : String(e); + toast(translateBackendError(t, message), 5000); + } +} + function formatMentionTableType(tableType: string) { const normalized = tableType.toUpperCase().replace(/\s+/g, "_"); if (normalized.includes("VIEW")) return t("ai.tableMentionTypes.view"); @@ -889,9 +1238,27 @@ function onPromptKeydown(event: KeyboardEvent) { } } +async function loadReferencedSqlFiles(mentions: AiSqlFileMention[]): Promise { + if (!mentions.length) return []; + const results: AiSqlFileContext[] = []; + for (const mention of mentions) { + const file = await savedSqlStore.ensureFileContent(mention.id).catch(() => undefined); + if (!file) continue; + const sql = file.sql || ""; + const truncated = sql.length > AI_SQL_FILE_CONTEXT_MAX_CHARS; + results.push({ + id: file.id, + name: file.name, + sql: truncated ? `${sql.slice(0, AI_SQL_FILE_CONTEXT_MAX_CHARS)}\n-- ... truncated ...` : sql, + truncated, + }); + } + return results; +} + async function send() { const text = prompt.value.trim(); - if ((!text && !selectedMentions.value.length) || isGenerating.value) return; + if ((!text && !selectedMentions.value.length && !selectedSqlFileMentions.value.length) || isGenerating.value) return; if (!props.connection || !props.tab) return; if (!settings.isConfigured()) { @@ -899,20 +1266,23 @@ async function send() { return; } - const mentionedTables = [...selectedMentions.value, ...parseAiTableMentions(text)]; - const displayText = [selectedMentions.value.map((mention) => mention.raw).join(" "), text].filter(Boolean).join(" "); + const selectedTableMentions = [...selectedMentions.value]; + const selectedSqlFiles = [...selectedSqlFileMentions.value]; + const mentionedTables = [...selectedTableMentions, ...parseAiTableMentions(text)]; + const modelInstruction = [selectedTableMentions.map((mention) => mention.raw).join(" "), selectedSqlFiles.map((mention) => mention.raw).join(" "), text].filter(Boolean).join(" "); - messages.value.push({ role: "user", content: displayText }); + messages.value.push({ role: "user", content: text, mentions: selectedMessageMentions(selectedTableMentions, selectedSqlFiles) }); // Save to prompt history (deduplicate consecutive duplicates) - if (displayText && promptHistory.value[0] !== displayText) { - promptHistory.value.unshift(displayText); + if (text && promptHistory.value[0] !== text) { + promptHistory.value.unshift(text); if (promptHistory.value.length > 100) promptHistory.value.length = 100; } historyIndex.value = -1; draftBeforeHistory.value = ""; prompt.value = ""; selectedMentions.value = []; - scrollToBottom(); + selectedSqlFileMentions.value = []; + scrollToBottom({ force: true }); const requestedAction = activeAction.value; const requestedMode = assistantMode.value; @@ -924,8 +1294,10 @@ async function send() { const agentEvents: AgentEvent[] = []; agentTokens.value = null; try { + const sqlFiles = await loadReferencedSqlFiles(selectedSqlFiles); const context = await buildAiContext(props.tab, props.connection, { mentionedTables, + sqlFiles, }); const history: AiMessage[] = messagesForAgentHistory(messages.value.slice(0, -2)); await runAgentStream( @@ -933,7 +1305,7 @@ async function send() { config: settings.aiConfig, action: activeAction.value, mode: requestedMode, - instruction: displayText, + instruction: modelInstruction, context, }, history, @@ -993,7 +1365,7 @@ async function send() { const agentPlan = buildAiAgentPlan({ mode: requestedMode, action: requestedAction, - instruction: displayText, + instruction: modelInstruction, assistantContent: msg?.content || "", connection: props.connection, }); @@ -1064,12 +1436,13 @@ async function persistConversation() { const first = messages.value.find((m) => m.role === "user" && m.kind !== "contextSummary"); await saveAiConversation({ id: conversationId.value, - title: first ? first.content.slice(0, 50) : "Untitled", + title: first ? messageTitle(first).slice(0, 50) : "Untitled", connectionName: props.connection.name, database: props.tab?.database || "", messages: messages.value.map((m) => ({ role: m.role, content: m.content, + ...(m.mentions?.length ? { mentions: m.mentions } : {}), ...(m.reasoning ? { reasoning: m.reasoning } : {}), ...(m.kind ? { kind: m.kind } : {}), })), @@ -1088,13 +1461,14 @@ function selectConversation(conv: AiConversation) { messages.value = conv.messages.map((m) => ({ role: m.role as "user" | "assistant", content: m.content, + mentions: Array.isArray(m.mentions) ? (m.mentions as AiMessageMention[]) : undefined, reasoning: m.reasoning, kind: m.kind, })); agentTokens.value = null; pendingCompaction.value = null; showConversationList.value = false; - scrollToBottom(); + scrollToBottom({ force: true }); } async function deleteConversation(id: string) { @@ -1185,6 +1559,7 @@ function stopResize() { onUnmounted(() => { clearTimeout(mentionTimer); cancelStream(); + detachMessageScrollListener(); // 清理拖拽事件监听，防止内存泄漏 document.removeEventListener("mousemove", handleResize); document.removeEventListener("mouseup", stopResize); @@ -1282,133 +1657,174 @@ async function openExternalUrl(url: string) {

{{ t("ai.welcome") }}

- -
-