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
@@ -4,6 +4,7 @@ import { uuid } from "@/lib/common/utils";
import { useI18n } from "vue-i18n"; import { useI18n } from "vue-i18n";
import { translateBackendError } from "@/i18n/backend-errors"; import { translateBackendError } from "@/i18n/backend-errors";
import { import {
ArrowDown,
ArrowUp, ArrowUp,
ArrowRightLeft, ArrowRightLeft,
AlertTriangle, AlertTriangle,
@@ -13,6 +14,7 @@ import {
CircleSlash, CircleSlash,
Copy, Copy,
Database, Database,
FileCode,
GitBranch, GitBranch,
HelpCircle, HelpCircle,
History, History,
@@ -44,11 +46,13 @@ import { useTheme } from "@/composables/useTheme";
import { useSettingsStore, AI_PROVIDER_PRESETS, type AiProvider } from "@/stores/settingsStore"; import { useSettingsStore, AI_PROVIDER_PRESETS, type AiProvider } from "@/stores/settingsStore";
import AiProviderLogo from "@/components/icons/AiProviderLogo.vue"; import AiProviderLogo from "@/components/icons/AiProviderLogo.vue";
import { useConnectionStore } from "@/stores/connectionStore"; import { useConnectionStore } from "@/stores/connectionStore";
import { useSavedSqlStore } from "@/stores/savedSqlStore";
import { connectionIconType } from "@/lib/connection/connectionPresentation"; import { connectionIconType } from "@/lib/connection/connectionPresentation";
import DatabaseIcon from "@/components/icons/DatabaseIcon.vue"; import DatabaseIcon from "@/components/icons/DatabaseIcon.vue";
import { useQueryStore } from "@/stores/queryStore"; import { useQueryStore } from "@/stores/queryStore";
import { useToast } from "@/composables/useToast"; 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 { formatAiModelOption } from "@/lib/ai/aiModelPresentation";
import type { AgentEvent } from "@/lib/backend/tauri"; import type { AgentEvent } from "@/lib/backend/tauri";
import { buildAiAgentPlan } from "@/lib/ai/aiAgentPlan"; import { buildAiAgentPlan } from "@/lib/ai/aiAgentPlan";
@@ -58,7 +62,7 @@ import { createAiMessageRenderer } from "@/lib/ai/aiMessageRender";
import { formatAiInlineMarkdown, handleAiMarkdownLinkClick } from "@/lib/ai/aiMarkdown"; import { formatAiInlineMarkdown, handleAiMarkdownLinkClick } from "@/lib/ai/aiMarkdown";
import { aiCancelStream, aiListModels, saveAiConversation, loadAiConversations, deleteAiConversation, listSchemas, listTables, type AiConversation, type AiModelInfo } from "@/lib/backend/api"; 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 { 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 { useDatabaseOptions } from "@/composables/useDatabaseOptions";
import { decodeSelectableDatabaseValue, encodeSelectableDatabaseValue, formatDatabaseLabel, resolveDefaultDatabase } from "@/lib/database/defaultDatabase"; import { decodeSelectableDatabaseValue, encodeSelectableDatabaseValue, formatDatabaseLabel, resolveDefaultDatabase } from "@/lib/database/defaultDatabase";
import { isSchemaAware } from "@/lib/database/databaseCapabilities"; import { isSchemaAware } from "@/lib/database/databaseCapabilities";
@@ -73,13 +77,36 @@ import { visibleToActualIndex } from "@/lib/ai/aiMessageEdit";
const { t } = useI18n(); const { t } = useI18n();
const settings = useSettingsStore(); const settings = useSettingsStore();
const connectionStore = useConnectionStore(); const connectionStore = useConnectionStore();
const savedSqlStore = useSavedSqlStore();
const queryStore = useQueryStore(); const queryStore = useQueryStore();
const { openTableTarget } = useNavigationTargets({
showFieldLineageDialog: ref(false),
showDatabaseSearchDialog: ref(false),
});
const { toast } = useToast(); const { toast } = useToast();
const { isDark } = useTheme(); 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 { interface ChatMessage {
role: "user" | "assistant"; role: "user" | "assistant";
content: string; content: string;
mentions?: AiMessageMention[];
reasoning?: string; reasoning?: string;
isThinking?: boolean; isThinking?: boolean;
agentSteps?: AiAgentStepItem[]; agentSteps?: AiAgentStepItem[];
@@ -111,6 +138,9 @@ const conversationId = ref("");
const conversations = ref<AiConversation[]>([]); const conversations = ref<AiConversation[]>([]);
const showConversationList = ref(false); const showConversationList = ref(false);
const promptTextareaRef = ref<HTMLTextAreaElement | null>(null); const promptTextareaRef = ref<HTMLTextAreaElement | null>(null);
const shouldAutoScroll = ref(true);
const userPausedAutoScroll = ref(false);
const showScrollToBottom = ref(false);
const promptCompositionActive = ref(false); const promptCompositionActive = ref(false);
const shikiCodeHighlighter = ref<AiCodeHighlighter>(); const shikiCodeHighlighter = ref<AiCodeHighlighter>();
const agentTokens = ref<{ input: number; output: number } | null>(null); const agentTokens = ref<{ input: number; output: number } | null>(null);
@@ -120,12 +150,21 @@ const draftBeforeHistory = ref("");
const editingMessageIndex = ref<number | null>(null); const editingMessageIndex = ref<number | null>(null);
const editingContent = ref(""); const editingContent = ref("");
const editingMentions = ref<AiPromptMentionChip[]>([]);
const editCompositionActive = ref(false); 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) { function startEditMessage(visibleIndex: number) {
if (isGenerating.value) return; if (isGenerating.value) return;
editingMessageIndex.value = visibleIndex; editingMessageIndex.value = visibleIndex;
editingContent.value = visibleMessages.value[visibleIndex].content; const msg = visibleMessages.value[visibleIndex];
editingContent.value = msg.content;
editingMentions.value = promptMentionChipsFromMessage(msg);
nextTick(() => { nextTick(() => {
const el = document.querySelector<HTMLTextAreaElement>("[data-edit-textarea]"); const el = document.querySelector<HTMLTextAreaElement>("[data-edit-textarea]");
if (el) { if (el) {
@@ -138,11 +177,12 @@ function startEditMessage(visibleIndex: number) {
function cancelEdit() { function cancelEdit() {
editingMessageIndex.value = null; editingMessageIndex.value = null;
editingContent.value = ""; editingContent.value = "";
editingMentions.value = [];
} }
function submitEdit(visibleIndex: number) { function submitEdit(visibleIndex: number) {
const content = editingContent.value.trim(); const content = editingContent.value.trim();
if (!content) return; if (!content && !editingMentions.value.length) return;
const actualIndex = visibleToActualIndex(messages.value, visibleIndex); const actualIndex = visibleToActualIndex(messages.value, visibleIndex);
if (actualIndex < 0) return; if (actualIndex < 0) return;
if (!props.connection || !props.tab) return; if (!props.connection || !props.tab) return;
@@ -153,7 +193,9 @@ function submitEdit(visibleIndex: number) {
messages.value = messages.value.slice(0, actualIndex); messages.value = messages.value.slice(0, actualIndex);
editingMessageIndex.value = null; editingMessageIndex.value = null;
editingContent.value = ""; 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; prompt.value = content;
send(); send();
} }
@@ -254,12 +296,31 @@ let resizeStartY = 0;
let resizeStartHeight = 0; let resizeStartHeight = 0;
let promptPanelResizeObserver: ResizeObserver | undefined; let promptPanelResizeObserver: ResizeObserver | undefined;
interface AiMentionCandidate { interface AiTableMentionCandidate {
kind: "table";
schema?: string; schema?: string;
name: string; name: string;
tableType: 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 mentionOpen = ref(false);
const mentionLoading = ref(false); const mentionLoading = ref(false);
const mentionError = ref(""); const mentionError = ref("");
@@ -269,9 +330,13 @@ const mentionCandidates = ref<AiMentionCandidate[]>([]);
const mentionCache = ref<Record<string, AiMentionCandidate[]>>({}); const mentionCache = ref<Record<string, AiMentionCandidate[]>>({});
const mentionListRef = ref<HTMLElement | null>(null); const mentionListRef = ref<HTMLElement | null>(null);
const selectedMentions = ref<AiTableMention[]>([]); const selectedMentions = ref<AiTableMention[]>([]);
const selectedSqlFileMentions = ref<AiSqlFileMention[]>([]);
let mentionTimer: ReturnType<typeof setTimeout> | undefined; let mentionTimer: ReturnType<typeof setTimeout> | undefined;
let mentionRequestId = 0; let mentionRequestId = 0;
const AI_SQL_FILE_MENTION_CANDIDATE_LIMIT = 50;
const AI_SQL_FILE_CONTEXT_MAX_CHARS = 12_000;
interface AiActionButton { interface AiActionButton {
action: AiAction; action: AiAction;
icon: Component; icon: Component;
@@ -348,22 +413,35 @@ function messagesForAgentHistory(historyMessages: ChatMessage[]): AiMessage[] {
} }
} }
if (latestSummaryIndex < 0) { 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 compactedHistory = historyMessages.slice(latestSummaryIndex);
const firstMsg = historyMessages[0]; const firstMsg = historyMessages[0];
if (firstMsg && firstMsg.role === "user" && firstMsg.kind !== "contextSummary") { 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 chatTitle = computed(() => {
const first = messages.value.find((m) => m.role === "user" && m.kind !== "contextSummary"); 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<AiPromptMentionChip[]>(() => [...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 isWaitingForFirstDelta = computed(() => {
const last = messages.value[messages.value.length - 1]; const last = messages.value[messages.value.length - 1];
@@ -511,7 +589,7 @@ function appendAssistantReasoning(assistantIdx: number, delta: string) {
scrollToBottom(); scrollToBottom();
} }
const reasoningExpanded = ref(false); const expandedReasoning = ref<Set<number>>(new Set());
const expandedSteps = ref<Set<string>>(new Set()); const expandedSteps = ref<Set<string>>(new Set());
function toggleStep(key: string) { function toggleStep(key: string) {
@@ -626,23 +704,169 @@ function agentEventToStep(event: AgentEvent, index: number): AiAgentStepItem | u
}; };
} }
function toggleReasoning() { function toggleReasoning(index: number) {
reasoningExpanded.value = !reasoningExpanded.value; 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 {
nextTick(() => {
const root = scrollRef.value?.$el as HTMLElement | undefined; const root = scrollRef.value?.$el as HTMLElement | undefined;
const el = root?.querySelector('[data-slot="scroll-area-viewport"]') as HTMLElement | null; 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 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; if (!el) return;
requestAnimationFrame(() => { requestAnimationFrame(() => {
if (!options.force && (userPausedAutoScroll.value || !shouldAutoScroll.value)) {
updateMessageScrollButtonVisibility();
return;
}
el.scrollTop = el.scrollHeight; 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) { 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[] { function mentionSchemaOrder(schemas: string[]): string[] {
@@ -688,10 +912,12 @@ async function loadMentionCandidates(query: string) {
mentionLoading.value = true; mentionLoading.value = true;
mentionError.value = ""; mentionError.value = "";
const { schemaPrefix, tableFilter } = normalizeMentionQuery(query); const { schemaPrefix, tableFilter } = normalizeMentionQuery(query);
let sqlFileCandidates: AiSqlFileMentionCandidate[] = [];
try { try {
sqlFileCandidates = await loadSqlFileMentionCandidates(query);
await connectionStore.ensureConnected(props.tab.connectionId); await connectionStore.ensureConnected(props.tab.connectionId);
let candidates: AiMentionCandidate[] = []; let tableCandidates: AiMentionCandidate[] = [];
if (isSchemaAware(props.connection.db_type)) { if (isSchemaAware(props.connection.db_type)) {
const schemas = mentionSchemaOrder(await listSchemas(props.tab.connectionId, props.tab.database)); const schemas = mentionSchemaOrder(await listSchemas(props.tab.connectionId, props.tab.database));
const filteredSchemas = schemaPrefix ? schemas.filter((schema) => schema.toLowerCase().includes(schemaPrefix.toLowerCase())) : schemas; 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 { } else {
const schema = props.tab.database || props.connection.database || "main"; 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); 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)), tables.map((table) => mentionCandidateFromTable(table)),
tableFilter, tableFilter,
AI_TABLE_MENTION_CANDIDATE_LIMIT, AI_TABLE_MENTION_CANDIDATE_LIMIT,
@@ -717,11 +943,18 @@ async function loadMentionCandidates(query: string) {
} }
if (requestId !== mentionRequestId) return; 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]; mentionCandidates.value = mentionCache.value[key];
setMentionSelectedIndex(0); setMentionSelectedIndex(0);
} catch (e: unknown) { } catch (e: unknown) {
if (requestId !== mentionRequestId) return; 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); const message = e instanceof Error ? e.message : String(e);
mentionError.value = translateBackendError(t, message); mentionError.value = translateBackendError(t, message);
mentionCandidates.value = []; mentionCandidates.value = [];
@@ -730,26 +963,142 @@ async function loadMentionCandidates(query: string) {
} }
} }
function mentionCandidateFromTable(table: TableInfo, schema?: string): AiMentionCandidate { async function loadSqlFileMentionCandidates(query: string): Promise<AiSqlFileMentionCandidate[]> {
return { schema, name: table.name, tableType: table.table_type }; 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("."); return [mention.schema, mention.table].filter(Boolean).join(".");
} }
function removeMentionChip(mention: AiTableMention) { 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); selectedMentions.value = selectedMentions.value.filter((item) => item.raw !== mention.raw);
}
nextTick(() => promptTextareaRef.value?.focus()); nextTick(() => promptTextareaRef.value?.focus());
} }
function removeEditingMentionChip(index: number) {
editingMentions.value = editingMentions.value.filter((_, itemIndex) => itemIndex !== index);
nextTick(() => {
const el = document.querySelector<HTMLTextAreaElement>("[data-edit-textarea]");
el?.focus();
});
}
function addSelectedMention(candidate: AiMentionCandidate) { 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 raw = formatAiTableMention(candidate.schema, candidate.name);
const key = `${candidate.schema || ""}.${candidate.name}`.toLowerCase(); const key = `${candidate.schema || ""}.${candidate.name}`.toLowerCase();
if (selectedMentions.value.some((mention) => `${mention.schema || ""}.${mention.table}`.toLowerCase() === key)) return; if (selectedMentions.value.some((mention) => `${mention.schema || ""}.${mention.table}`.toLowerCase() === key)) return;
selectedMentions.value.push({ raw, schema: candidate.schema, table: candidate.name }); 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) { function formatMentionTableType(tableType: string) {
const normalized = tableType.toUpperCase().replace(/\s+/g, "_"); const normalized = tableType.toUpperCase().replace(/\s+/g, "_");
if (normalized.includes("VIEW")) return t("ai.tableMentionTypes.view"); if (normalized.includes("VIEW")) return t("ai.tableMentionTypes.view");
@@ -889,9 +1238,27 @@ function onPromptKeydown(event: KeyboardEvent) {
} }
} }
async function loadReferencedSqlFiles(mentions: AiSqlFileMention[]): Promise<AiSqlFileContext[]> {
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() { async function send() {
const text = prompt.value.trim(); 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 (!props.connection || !props.tab) return;
if (!settings.isConfigured()) { if (!settings.isConfigured()) {
@@ -899,20 +1266,23 @@ async function send() {
return; return;
} }
const mentionedTables = [...selectedMentions.value, ...parseAiTableMentions(text)]; const selectedTableMentions = [...selectedMentions.value];
const displayText = [selectedMentions.value.map((mention) => mention.raw).join(" "), text].filter(Boolean).join(" "); 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) // Save to prompt history (deduplicate consecutive duplicates)
if (displayText && promptHistory.value[0] !== displayText) { if (text && promptHistory.value[0] !== text) {
promptHistory.value.unshift(displayText); promptHistory.value.unshift(text);
if (promptHistory.value.length > 100) promptHistory.value.length = 100; if (promptHistory.value.length > 100) promptHistory.value.length = 100;
} }
historyIndex.value = -1; historyIndex.value = -1;
draftBeforeHistory.value = ""; draftBeforeHistory.value = "";
prompt.value = ""; prompt.value = "";
selectedMentions.value = []; selectedMentions.value = [];
scrollToBottom(); selectedSqlFileMentions.value = [];
scrollToBottom({ force: true });
const requestedAction = activeAction.value; const requestedAction = activeAction.value;
const requestedMode = assistantMode.value; const requestedMode = assistantMode.value;
@@ -924,8 +1294,10 @@ async function send() {
const agentEvents: AgentEvent[] = []; const agentEvents: AgentEvent[] = [];
agentTokens.value = null; agentTokens.value = null;
try { try {
const sqlFiles = await loadReferencedSqlFiles(selectedSqlFiles);
const context = await buildAiContext(props.tab, props.connection, { const context = await buildAiContext(props.tab, props.connection, {
mentionedTables, mentionedTables,
sqlFiles,
}); });
const history: AiMessage[] = messagesForAgentHistory(messages.value.slice(0, -2)); const history: AiMessage[] = messagesForAgentHistory(messages.value.slice(0, -2));
await runAgentStream( await runAgentStream(
@@ -933,7 +1305,7 @@ async function send() {
config: settings.aiConfig, config: settings.aiConfig,
action: activeAction.value, action: activeAction.value,
mode: requestedMode, mode: requestedMode,
instruction: displayText, instruction: modelInstruction,
context, context,
}, },
history, history,
@@ -993,7 +1365,7 @@ async function send() {
const agentPlan = buildAiAgentPlan({ const agentPlan = buildAiAgentPlan({
mode: requestedMode, mode: requestedMode,
action: requestedAction, action: requestedAction,
instruction: displayText, instruction: modelInstruction,
assistantContent: msg?.content || "", assistantContent: msg?.content || "",
connection: props.connection, connection: props.connection,
}); });
@@ -1064,12 +1436,13 @@ async function persistConversation() {
const first = messages.value.find((m) => m.role === "user" && m.kind !== "contextSummary"); const first = messages.value.find((m) => m.role === "user" && m.kind !== "contextSummary");
await saveAiConversation({ await saveAiConversation({
id: conversationId.value, id: conversationId.value,
title: first ? first.content.slice(0, 50) : "Untitled", title: first ? messageTitle(first).slice(0, 50) : "Untitled",
connectionName: props.connection.name, connectionName: props.connection.name,
database: props.tab?.database || "", database: props.tab?.database || "",
messages: messages.value.map((m) => ({ messages: messages.value.map((m) => ({
role: m.role, role: m.role,
content: m.content, content: m.content,
...(m.mentions?.length ? { mentions: m.mentions } : {}),
...(m.reasoning ? { reasoning: m.reasoning } : {}), ...(m.reasoning ? { reasoning: m.reasoning } : {}),
...(m.kind ? { kind: m.kind } : {}), ...(m.kind ? { kind: m.kind } : {}),
})), })),
@@ -1088,13 +1461,14 @@ function selectConversation(conv: AiConversation) {
messages.value = conv.messages.map((m) => ({ messages.value = conv.messages.map((m) => ({
role: m.role as "user" | "assistant", role: m.role as "user" | "assistant",
content: m.content, content: m.content,
mentions: Array.isArray(m.mentions) ? (m.mentions as AiMessageMention[]) : undefined,
reasoning: m.reasoning, reasoning: m.reasoning,
kind: m.kind, kind: m.kind,
})); }));
agentTokens.value = null; agentTokens.value = null;
pendingCompaction.value = null; pendingCompaction.value = null;
showConversationList.value = false; showConversationList.value = false;
scrollToBottom(); scrollToBottom({ force: true });
} }
async function deleteConversation(id: string) { async function deleteConversation(id: string) {
@@ -1185,6 +1559,7 @@ function stopResize() {
onUnmounted(() => { onUnmounted(() => {
clearTimeout(mentionTimer); clearTimeout(mentionTimer);
cancelStream(); cancelStream();
detachMessageScrollListener();
// 清理拖拽事件监听,防止内存泄漏 // 清理拖拽事件监听,防止内存泄漏
document.removeEventListener("mousemove", handleResize); document.removeEventListener("mousemove", handleResize);
document.removeEventListener("mouseup", stopResize); document.removeEventListener("mouseup", stopResize);
@@ -1282,12 +1657,28 @@ async function openExternalUrl(url: string) {
<Bot class="h-10 w-10 mb-3 opacity-30" /> <Bot class="h-10 w-10 mb-3 opacity-30" />
<p class="text-sm">{{ t("ai.welcome") }}</p> <p class="text-sm">{{ t("ai.welcome") }}</p>
</div> </div>
<ScrollArea v-else ref="scrollRef" class="min-h-0 flex-1 overflow-hidden"> <div v-else class="relative min-h-0 flex-1">
<ScrollArea ref="scrollRef" class="ai-message-scroll h-full overflow-hidden">
<div class="flex flex-col gap-3 p-3"> <div class="flex flex-col gap-3 p-3">
<template v-for="(msg, i) in visibleMessages" :key="i"> <template v-for="(msg, i) in visibleMessages" :key="i">
<div v-if="msg.role === 'user'" class="group flex justify-end"> <div v-if="msg.role === 'user'" class="group flex justify-end">
<div class="max-w-[85%]"> <div class="max-w-[85%]">
<template v-if="editingMessageIndex === i"> <template v-if="editingMessageIndex === i">
<div v-if="editingMentions.length" class="mb-1.5 flex flex-wrap justify-end gap-1">
<button
v-for="(mention, mentionIndex) in editingMentions"
:key="`${mention.kind}:${mention.raw}:${mentionIndex}`"
type="button"
class="group inline-flex max-w-full items-center gap-1 rounded border border-border/80 bg-muted/70 px-1.5 py-0.5 text-[11px] text-foreground/90 hover:bg-muted"
:title="mentionDisplayName(mention)"
@click="removeEditingMentionChip(mentionIndex)"
>
<FileCode v-if="mention.kind === 'sqlFile'" class="h-3 w-3 shrink-0 text-primary" />
<Table2 v-else class="h-3 w-3 shrink-0 text-primary" />
<span class="truncate">{{ mentionDisplayName(mention) }}</span>
<X class="h-3 w-3 shrink-0 text-muted-foreground group-hover:text-foreground" />
</button>
</div>
<textarea <textarea
data-edit-textarea data-edit-textarea
v-model="editingContent" v-model="editingContent"
@@ -1307,8 +1698,22 @@ async function openExternalUrl(url: string) {
<button v-if="!isGenerating" class="mt-1 hidden h-5 w-5 shrink-0 items-center justify-center rounded text-muted-foreground hover:bg-muted hover:text-foreground group-hover:flex" :title="t('ai.editMessage')" @click="startEditMessage(i)"> <button v-if="!isGenerating" class="mt-1 hidden h-5 w-5 shrink-0 items-center justify-center rounded text-muted-foreground hover:bg-muted hover:text-foreground group-hover:flex" :title="t('ai.editMessage')" @click="startEditMessage(i)">
<Pencil class="h-3 w-3" /> <Pencil class="h-3 w-3" />
</button> </button>
<div class="whitespace-pre-wrap rounded-lg bg-primary px-3 py-2 text-xs text-primary-foreground"> <div class="min-w-0 rounded-lg bg-primary px-3 py-2 text-xs text-primary-foreground">
{{ msg.content }} <div v-if="msg.mentions?.length" class="mb-1.5 flex flex-wrap justify-end gap-1">
<button
v-for="mention in msg.mentions"
:key="`${mention.kind}:${mention.raw}`"
type="button"
class="inline-flex max-w-full items-center gap-1 rounded border border-primary-foreground/25 bg-primary-foreground/15 px-1.5 py-0.5 text-[11px] text-primary-foreground hover:bg-primary-foreground/25"
:title="mention.kind === 'sqlFile' ? mention.name : [mention.schema, mention.table].filter(Boolean).join('.')"
@click.stop="openMessageMention(mention)"
>
<FileCode v-if="mention.kind === 'sqlFile'" class="h-3 w-3 shrink-0" />
<Table2 v-else class="h-3 w-3 shrink-0" />
<span class="truncate">{{ mention.kind === "sqlFile" ? mention.name : [mention.schema, mention.table].filter(Boolean).join(".") }}</span>
</button>
</div>
<div v-if="msg.content" class="whitespace-pre-wrap">{{ msg.content }}</div>
</div> </div>
</div> </div>
</template> </template>
@@ -1318,16 +1723,16 @@ async function openExternalUrl(url: string) {
<div v-else-if="msg.content || msg.reasoning || msg.isThinking" class="flex"> <div v-else-if="msg.content || msg.reasoning || msg.isThinking" class="flex">
<div class="max-w-[95%] min-w-0 rounded-lg bg-muted px-3 py-2 text-xs leading-relaxed"> <div class="max-w-[95%] min-w-0 rounded-lg bg-muted px-3 py-2 text-xs leading-relaxed">
<div v-if="msg.reasoning || msg.isThinking" class="mb-2"> <div v-if="msg.reasoning || msg.isThinking" class="mb-2">
<button class="flex items-center gap-1 text-[11px] text-muted-foreground hover:text-foreground transition-colors" @click="toggleReasoning()"> <button class="flex items-center gap-1 text-[11px] text-muted-foreground hover:text-foreground transition-colors" @click="toggleReasoning(i)">
<ChevronRight class="h-3 w-3 transition-transform duration-200" :class="{ 'rotate-90': reasoningExpanded }" /> <ChevronRight class="h-3 w-3 transition-transform duration-200" :class="{ 'rotate-90': expandedReasoning.has(i) || msg.isThinking }" />
<Loader2 v-if="msg.isThinking" class="h-3 w-3 animate-spin" /> <Loader2 v-if="msg.isThinking" class="h-3 w-3 animate-spin" />
<span>{{ t("ai.reasoningProcess") }}</span> <span>{{ t("ai.reasoningProcess") }}</span>
</button> </button>
<div <div
class="overflow-hidden transition-[max-height,opacity] duration-200 ease-in-out" class="overflow-hidden transition-[max-height,opacity] duration-200 ease-in-out"
:style="{ :style="{
maxHeight: reasoningExpanded ? '20000px' : '0px', maxHeight: expandedReasoning.has(i) || msg.isThinking ? '20000px' : '0px',
opacity: reasoningExpanded ? '1' : '0', opacity: expandedReasoning.has(i) || msg.isThinking ? '1' : '0',
}" }"
> >
<div class="mt-1.5 pl-4 border-l-2 border-muted-foreground/20 text-[11px] text-muted-foreground whitespace-pre-wrap"> <div class="mt-1.5 pl-4 border-l-2 border-muted-foreground/20 text-[11px] text-muted-foreground whitespace-pre-wrap">
@@ -1409,6 +1814,17 @@ async function openExternalUrl(url: string) {
</div> </div>
</div> </div>
</ScrollArea> </ScrollArea>
<button
v-if="showScrollToBottom"
type="button"
class="absolute bottom-3 right-3 z-10 inline-flex h-8 w-8 items-center justify-center rounded-full border bg-background/95 text-foreground shadow-md backdrop-blur hover:bg-muted"
:title="t('ai.scrollToBottom')"
@click="scrollToBottom({ force: true })"
>
<ArrowDown class="h-4 w-4" />
<span class="sr-only">{{ t("ai.scrollToBottom") }}</span>
</button>
</div>
<div class="p-2"> <div class="p-2">
<div ref="promptPanelRef" class="relative rounded-[6px] border bg-background"> <div ref="promptPanelRef" class="relative rounded-[6px] border bg-background">
@@ -1476,7 +1892,7 @@ async function openExternalUrl(url: string) {
<div v-else ref="mentionListRef" class="max-h-56 overflow-auto p-1"> <div v-else ref="mentionListRef" class="max-h-56 overflow-auto p-1">
<button <button
v-for="(candidate, index) in mentionCandidates" v-for="(candidate, index) in mentionCandidates"
:key="`${candidate.schema || ''}.${candidate.name}`" :key="candidate.kind === 'sqlFile' ? `sql-file:${candidate.id}` : `table:${candidate.schema || ''}.${candidate.name}`"
type="button" type="button"
:data-mention-index="index" :data-mention-index="index"
class="flex w-full min-w-0 items-center gap-2 rounded px-2 py-1.5 text-left text-xs hover:bg-muted" class="flex w-full min-w-0 items-center gap-2 rounded px-2 py-1.5 text-left text-xs hover:bg-muted"
@@ -1484,11 +1900,12 @@ async function openExternalUrl(url: string) {
@mousedown.prevent="insertMention(candidate)" @mousedown.prevent="insertMention(candidate)"
@mouseenter="setMentionSelectedIndex(index, false)" @mouseenter="setMentionSelectedIndex(index, false)"
> >
<Table2 class="h-3.5 w-3.5 shrink-0 text-muted-foreground" /> <FileCode v-if="candidate.kind === 'sqlFile'" class="h-3.5 w-3.5 shrink-0 text-primary" />
<Table2 v-else class="h-3.5 w-3.5 shrink-0 text-muted-foreground" />
<span class="min-w-0 flex-1 truncate"> <span class="min-w-0 flex-1 truncate">
<template v-if="candidate.schema">{{ candidate.schema }}.</template>{{ candidate.name }} {{ mentionCandidateName(candidate) }}
</span> </span>
<span class="shrink-0 text-[10px] text-muted-foreground">{{ formatMentionTableType(candidate.tableType) }}</span> <span class="max-w-[45%] shrink-0 truncate text-[10px] text-muted-foreground">{{ formatMentionCandidateType(candidate) }}</span>
</button> </button>
</div> </div>
</div> </div>
@@ -1501,7 +1918,8 @@ async function openExternalUrl(url: string) {
:title="mentionDisplayName(mention)" :title="mentionDisplayName(mention)"
@click="removeMentionChip(mention)" @click="removeMentionChip(mention)"
> >
<Table2 class="h-3 w-3 shrink-0 text-primary" /> <FileCode v-if="mention.kind === 'sqlFile'" class="h-3 w-3 shrink-0 text-primary" />
<Table2 v-else class="h-3 w-3 shrink-0 text-primary" />
<span class="truncate">{{ mentionDisplayName(mention) }}</span> <span class="truncate">{{ mentionDisplayName(mention) }}</span>
<X class="h-3 w-3 shrink-0 text-muted-foreground group-hover:text-foreground" /> <X class="h-3 w-3 shrink-0 text-muted-foreground group-hover:text-foreground" />
</button> </button>
@@ -1594,7 +2012,7 @@ async function openExternalUrl(url: string) {
<button v-if="isGenerating" class="h-7 w-7 shrink-0 rounded-full bg-destructive text-destructive-foreground flex items-center justify-center" :title="t('ai.stopGenerating')" @click="cancelStream"> <button v-if="isGenerating" class="h-7 w-7 shrink-0 rounded-full bg-destructive text-destructive-foreground flex items-center justify-center" :title="t('ai.stopGenerating')" @click="cancelStream">
<Square class="h-3.5 w-3.5" /> <Square class="h-3.5 w-3.5" />
</button> </button>
<button v-else class="h-7 w-7 shrink-0 rounded-full bg-foreground text-background flex items-center justify-center disabled:opacity-30" :disabled="!prompt.trim() || !props.tab?.database" @click="send"> <button v-else class="h-7 w-7 shrink-0 rounded-full bg-foreground text-background flex items-center justify-center disabled:opacity-30" :disabled="(!prompt.trim() && !selectedMentions.length && !selectedSqlFileMentions.length) || !props.tab?.database" @click="send">
<ArrowUp class="h-4 w-4" /> <ArrowUp class="h-4 w-4" />
</button> </button>
</div> </div>
@@ -1705,6 +2123,10 @@ async function openExternalUrl(url: string) {
min-height: 1lh; min-height: 1lh;
} }
.ai-message-scroll :deep([data-slot="scroll-area-viewport"]) {
overflow-anchor: none;
}
.resize-handle { .resize-handle {
height: 4px; height: 4px;
width: 100%; width: 100%;
+3 -2
View File
@@ -1205,8 +1205,9 @@ export default {
thinking: "Thinking...", thinking: "Thinking...",
reasoningProcess: "Thinking process", reasoningProcess: "Thinking process",
stopGenerating: "Stop generating", stopGenerating: "Stop generating",
tableMentionEmpty: "No matching tables", scrollToBottom: "Back to bottom",
tableMentionPlaceholderHint: "Type {'@'} to mention tables for context", tableMentionEmpty: "No matching tables or SQL files",
tableMentionPlaceholderHint: "Type {'@'} to mention tables or SQL files for context",
tableMentionTypes: { tableMentionTypes: {
table: "Table", table: "Table",
view: "View", view: "View",
+3 -2
View File
@@ -1178,8 +1178,9 @@ export default withEnglishFallback({
thinking: "Procesando...", thinking: "Procesando...",
reasoningProcess: "Proceso de razonamiento", reasoningProcess: "Proceso de razonamiento",
stopGenerating: "Detener generación", stopGenerating: "Detener generación",
tableMentionEmpty: "Sin tablas coincidentes", scrollToBottom: "Volver al final",
tableMentionPlaceholderHint: "Escribe {'@'} para mencionar tablas como contexto", tableMentionEmpty: "Sin tablas ni archivos SQL coincidentes",
tableMentionPlaceholderHint: "Escribe {'@'} para mencionar tablas o archivos SQL como contexto",
tableMentionTypes: { tableMentionTypes: {
table: "Tabla", table: "Tabla",
view: "Vista", view: "Vista",
+3 -2
View File
@@ -1159,8 +1159,9 @@ export default withEnglishFallback({
thinking: "Riflessione...", thinking: "Riflessione...",
reasoningProcess: "Processo di riflessione", reasoningProcess: "Processo di riflessione",
stopGenerating: "Interrompi generazione", stopGenerating: "Interrompi generazione",
tableMentionEmpty: "Nessuna tabella corrispondente", scrollToBottom: "Torna in fondo",
tableMentionPlaceholderHint: "Digita {'@'} per menzionare le tabelle per il contesto", tableMentionEmpty: "Nessuna tabella o file SQL corrispondente",
tableMentionPlaceholderHint: "Digita {'@'} per menzionare tabelle o file SQL per il contesto",
tableMentionTypes: { tableMentionTypes: {
table: "Tabella", table: "Tabella",
view: "Vista", view: "Vista",
+3 -2
View File
@@ -1181,8 +1181,9 @@ export default withEnglishFallback({
thinking: "思考中...", thinking: "思考中...",
reasoningProcess: "思考プロセス", reasoningProcess: "思考プロセス",
stopGenerating: "生成を停止", stopGenerating: "生成を停止",
tableMentionEmpty: "一致するテーブルがありません", scrollToBottom: "一番下へ戻る",
tableMentionPlaceholderHint: "{'@'}を入力してコンテキスト用のテーブルを指定", tableMentionEmpty: "一致するテーブルまたは SQL ファイルがありません",
tableMentionPlaceholderHint: "{'@'}を入力してコンテキスト用のテーブルまたは SQL ファイルを指定",
tableMentionTypes: { tableMentionTypes: {
table: "テーブル", table: "テーブル",
view: "ビュー", view: "ビュー",
+3 -2
View File
@@ -1179,8 +1179,9 @@ export default withEnglishFallback({
thinking: "Pensando...", thinking: "Pensando...",
reasoningProcess: "Processo de raciocínio", reasoningProcess: "Processo de raciocínio",
stopGenerating: "Parar geração", stopGenerating: "Parar geração",
tableMentionEmpty: "Nenhuma tabela correspondente", scrollToBottom: "Voltar ao fim",
tableMentionPlaceholderHint: "Digite {'@'} para mencionar tabelas como contexto", tableMentionEmpty: "Nenhuma tabela ou arquivo SQL correspondente",
tableMentionPlaceholderHint: "Digite {'@'} para mencionar tabelas ou arquivos SQL como contexto",
tableMentionTypes: { tableMentionTypes: {
table: "Tabela", table: "Tabela",
view: "Visão", view: "Visão",
+3 -2
View File
@@ -1207,8 +1207,9 @@ export default withEnglishFallback({
thinking: "思考中...", thinking: "思考中...",
reasoningProcess: "思考过程", reasoningProcess: "思考过程",
stopGenerating: "停止生成", stopGenerating: "停止生成",
tableMentionEmpty: "没有匹配的表", scrollToBottom: "回到底部",
tableMentionPlaceholderHint: "输入 {'@'} 可选择表并带入上下文", tableMentionEmpty: "没有匹配的表或 SQL 文件",
tableMentionPlaceholderHint: "输入 {'@'} 可选择表或 SQL 文件并带入上下文",
tableMentionTypes: { tableMentionTypes: {
table: "表", table: "表",
view: "视图", view: "视图",
+3 -2
View File
@@ -1160,8 +1160,9 @@ export default withEnglishFallback({
thinking: "思考中……", thinking: "思考中……",
reasoningProcess: "思考過程", reasoningProcess: "思考過程",
stopGenerating: "停止產生", stopGenerating: "停止產生",
tableMentionEmpty: "沒有相符的資料表", scrollToBottom: "回到底部",
tableMentionPlaceholderHint: "輸入 {'@'} 可提及資料表作為脈絡", tableMentionEmpty: "沒有相符的資料表或 SQL 檔案",
tableMentionPlaceholderHint: "輸入 {'@'} 可提及資料表或 SQL 檔案作為脈絡",
tableMentionTypes: { tableMentionTypes: {
table: "資料表", table: "資料表",
view: "檢視", view: "檢視",
@@ -10,6 +10,7 @@ function context(overrides: Partial<AiContext> = {}): AiContext {
database: "app", database: "app",
currentSql: "", currentSql: "",
tables: [], tables: [],
sqlFiles: [],
truncated: false, truncated: false,
...overrides, ...overrides,
}; };
+28 -1
View File
@@ -66,6 +66,13 @@ export interface AiSchemaTable {
foreignKeys?: ForeignKeyInfo[]; foreignKeys?: ForeignKeyInfo[];
} }
export interface AiSqlFileContext {
id: string;
name: string;
sql: string;
truncated?: boolean;
}
export interface AiContext { export interface AiContext {
connectionId: string; connectionId: string;
connectionName: string; connectionName: string;
@@ -75,6 +82,7 @@ export interface AiContext {
lastError?: string; lastError?: string;
lastResultPreview?: string; lastResultPreview?: string;
tables: AiSchemaTable[]; tables: AiSchemaTable[];
sqlFiles: AiSqlFileContext[];
schemaScope?: "focused_table" | "database"; schemaScope?: "focused_table" | "database";
truncated: boolean; truncated: boolean;
} }
@@ -196,6 +204,7 @@ export function buildSystemPrompt(action: AiAction, context: AiContext, mode: Ai
const schema = formatSchema(context); const schema = formatSchema(context);
const resultPreview = context.lastResultPreview ? `\nLast result preview:\n${context.lastResultPreview}\n` : ""; const resultPreview = context.lastResultPreview ? `\nLast result preview:\n${context.lastResultPreview}\n` : "";
const lastError = context.lastError ? `\nLast error:\n${context.lastError}\n` : ""; const lastError = context.lastError ? `\nLast error:\n${context.lastError}\n` : "";
const referencedSqlFiles = formatReferencedSqlFiles(context);
const schemaScope = context.schemaScope ?? "database"; const schemaScope = context.schemaScope ?? "database";
const isZh = isChineseLocale(currentLocale()); const isZh = isChineseLocale(currentLocale());
@@ -225,6 +234,7 @@ export function buildSystemPrompt(action: AiAction, context: AiContext, mode: Ai
schemaCoverageLine(context, isZh), schemaCoverageLine(context, isZh),
"", "",
`Current SQL:\n${context.currentSql.trim() || "(empty)"}`, `Current SQL:\n${context.currentSql.trim() || "(empty)"}`,
referencedSqlFiles,
lastError, lastError,
resultPreview, resultPreview,
`Schema:\n${schema}`, `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 ? "当用户要求分析或查看某个表时,生成 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 中不存在的表或列。" : "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 ? "用户输入中的 @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 ? "不要生成多语句 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 ? "对于 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.", 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 schema = formatSchema(context);
const resultPreview = context.lastResultPreview ? `\nLast result preview:\n${context.lastResultPreview}\n` : ""; const resultPreview = context.lastResultPreview ? `\nLast result preview:\n${context.lastResultPreview}\n` : "";
const lastError = context.lastError ? `\nLast error:\n${context.lastError}\n` : ""; const lastError = context.lastError ? `\nLast error:\n${context.lastError}\n` : "";
const referencedSqlFiles = formatReferencedSqlFiles(context);
const lines: string[] = [ const lines: string[] = [
isZh ? `你是 DBX 内置的向量数据库助手。当前连接的是 ${dbLabel(context.databaseType)} 数据库。用中文回复。` : `You are DBX's vector database assistant. Connected to ${dbLabel(context.databaseType)}. Reply in English.`, 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.", 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), schemaCoverageLine(context, isZh),
"", "",
`Current collection:\n${context.currentSql.trim() || "(none)"}`, `Current collection:\n${context.currentSql.trim() || "(none)"}`,
referencedSqlFiles,
lastError, lastError,
resultPreview, resultPreview,
"", "",
@@ -369,7 +382,20 @@ function formatSchema(context: AiContext): string {
.join("\n\n"); .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 maxTables = options.maxTables ?? 50;
const maxColumnsPerTable = options.maxColumnsPerTable ?? 40; const maxColumnsPerTable = options.maxColumnsPerTable ?? 40;
const maxIndexesPerTable = options.maxIndexesPerTable ?? 10; const maxIndexesPerTable = options.maxIndexesPerTable ?? 10;
@@ -496,6 +522,7 @@ export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig
lastError: extractLastError(tab.result), lastError: extractLastError(tab.result),
lastResultPreview: formatResultPreview(tab.result), lastResultPreview: formatResultPreview(tab.result),
tables, tables,
sqlFiles: options.sqlFiles ?? [],
schemaScope, schemaScope,
truncated, truncated,
}; };
+1
View File
@@ -538,6 +538,7 @@ export async function listSqlFilesInFolder(folderPath: string): Promise<SqlFileE
export interface AiChatMessage { export interface AiChatMessage {
role: string; role: string;
content: string; content: string;
mentions?: unknown[];
reasoning?: string; reasoning?: string;
kind?: "contextSummary"; kind?: "contextSummary";
} }
+2
View File
@@ -194,6 +194,8 @@ pub struct AiChatMessage {
pub role: String, pub role: String,
pub content: String, pub content: String,
#[serde(default, skip_serializing_if = "Option::is_none")] #[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>, pub reasoning: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub kind: Option<String>, pub kind: Option<String>,
+25
View File
@@ -34,6 +34,7 @@ const { buildSystemPrompt, isVectorDbType, buildUserPrompt } = await import("../
function context(overrides: Partial<AiContext> = {}): AiContext { function context(overrides: Partial<AiContext> = {}): AiContext {
return { return {
connectionId: "conn-1",
connectionName: "prod-analytics", connectionName: "prod-analytics",
databaseType: "postgres", databaseType: "postgres",
database: "app", database: "app",
@@ -52,6 +53,7 @@ function context(overrides: Partial<AiContext> = {}): AiContext {
foreignKeys: [{ column: "user_id", ref_table: "users", ref_column: "id" }], foreignKeys: [{ column: "user_id", ref_table: "users", ref_column: "id" }],
}, },
], ],
sqlFiles: [],
truncated: false, truncated: false,
...overrides, ...overrides,
}; };
@@ -99,10 +101,32 @@ test("prompt enforces database dialect and single executable statement safety",
assert.match(prompt, /不要在同一个回答里混合 SELECT 和写操作/); 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 // Vector database tests
function vectorContext(overrides: Partial<AiContext> = {}): AiContext { function vectorContext(overrides: Partial<AiContext> = {}): AiContext {
return { return {
connectionId: "conn-1",
connectionName: "my-qdrant", connectionName: "my-qdrant",
databaseType: "qdrant", databaseType: "qdrant",
database: "default", database: "default",
@@ -115,6 +139,7 @@ function vectorContext(overrides: Partial<AiContext> = {}): AiContext {
columns: [], columns: [],
}, },
], ],
sqlFiles: [],
truncated: false, truncated: false,
...overrides, ...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 { function baseContext(overrides: Partial<AiContext> = {}): AiContext {
return { return {
connectionId: "conn-1",
connectionName: "prod-analytics", connectionName: "prod-analytics",
databaseType: "postgres", databaseType: "postgres",
database: "warehouse", database: "warehouse",
@@ -67,6 +68,7 @@ function baseContext(overrides: Partial<AiContext> = {}): AiContext {
], ],
}, },
], ],
sqlFiles: [],
truncated: false, truncated: false,
...overrides, ...overrides,
}; };
@@ -4,6 +4,7 @@ import { test } from "vitest";
import { buildSystemPrompt, type AiContext } from "../../apps/desktop/src/lib/ai/ai.ts"; import { buildSystemPrompt, type AiContext } from "../../apps/desktop/src/lib/ai/ai.ts";
const BASE_CONTEXT: AiContext = { const BASE_CONTEXT: AiContext = {
connectionId: "conn-1",
connectionName: "local", connectionName: "local",
databaseType: "mysql", databaseType: "mysql",
database: "app", database: "app",
@@ -35,6 +36,7 @@ const BASE_CONTEXT: AiContext = {
], ],
}, },
], ],
sqlFiles: [],
truncated: false, truncated: false,
}; };