diff --git a/package-lock.json b/package-lock.json index cff4e1a295e..8f54dd1b5c5 100644 --- a/package-lock.json +++ b/package-lock.json @@ -39064,6 +39064,7 @@ "clipboard": "^2.0.8", "clsx": "^1.2.1", "common": "*", + "common-tags": "^1.8.2", "config": "*", "configcat-js": "^7.0.0", "dayjs": "^1.11.0", @@ -39138,6 +39139,7 @@ "@testing-library/react-hooks": "^8.0.1", "@testing-library/user-event": "^13.5.0", "@types/blueimp-md5": "^2.18.0", + "@types/common-tags": "^1.8.1", "@types/file-saver": "^2.0.2", "@types/json-logic-js": "^1.2.1", "@types/lodash": "^4.14.172", @@ -66062,6 +66064,7 @@ "@testing-library/react-hooks": "^8.0.1", "@testing-library/user-event": "^13.5.0", "@types/blueimp-md5": "^2.18.0", + "@types/common-tags": "^1.8.1", "@types/file-saver": "^2.0.2", "@types/json-logic-js": "^1.2.1", "@types/lodash": "^4.14.172", @@ -66092,6 +66095,7 @@ "clipboard": "^2.0.8", "clsx": "^1.2.1", "common": "*", + "common-tags": "^1.8.2", "config": "*", "configcat-js": "^7.0.0", "dayjs": "^1.11.0", diff --git a/packages/ui/src/components/Command/AiCommand.tsx b/packages/ui/src/components/Command/AiCommand.tsx index 78314b8858e..5e816dd7d53 100644 --- a/packages/ui/src/components/Command/AiCommand.tsx +++ b/packages/ui/src/components/Command/AiCommand.tsx @@ -272,6 +272,70 @@ export function useAiChat({ } } +/** + * Perform a one-off query to AI based on a snapshot of messages + */ +export function queryAi(messages: Message[], timeout = 0) { + return new Promise((resolve, reject) => { + const eventSource = new SSE(`${edgeFunctionUrl}/ai-docs`, { + headers: { + apikey: process.env.NEXT_PUBLIC_SUPABASE_ANON_KEY ?? '', + Authorization: `Bearer ${process.env.NEXT_PUBLIC_SUPABASE_ANON_KEY}`, + 'Content-Type': 'application/json', + }, + payload: JSON.stringify({ + messages: messages.map(({ role, content }) => ({ role, content })), + }), + }) + + let timeoutId: number | undefined + + function handleError(err: T) { + if (timeoutId) { + clearTimeout(timeoutId) + } + console.error(err) + reject(err) + } + + if (timeout > 0) { + timeoutId = window.setTimeout(() => { + handleError(new Error('AI query timed out')) + }, timeout) + } + + let answer = '' + + eventSource.addEventListener('error', handleError) + eventSource.addEventListener('message', (e) => { + try { + if (e.data === '[DONE]') { + if (timeoutId) { + clearTimeout(timeoutId) + } + resolve(answer) + return + } + + const completionResponse: CreateChatCompletionResponse = JSON.parse(e.data) + const [ + { + delta: { content }, + }, + ] = completionResponse.choices as CreateChatCompletionResponseChoicesInnerDelta[] + + if (content) { + answer += content + } + } catch (err) { + handleError(err) + } + }) + + eventSource.stream() + }) +} + const AiCommand = () => { const { isLoading, setIsLoading, search, setSearch, MarkdownHandler } = useCommandMenu() diff --git a/packages/ui/src/components/Command/CommandMenuProvider.tsx b/packages/ui/src/components/Command/CommandMenuProvider.tsx index 3024ee15991..7ebd7e9fb2d 100644 --- a/packages/ui/src/components/Command/CommandMenuProvider.tsx +++ b/packages/ui/src/components/Command/CommandMenuProvider.tsx @@ -25,7 +25,7 @@ export interface CommandMenuContextValue { MarkdownHandler: (props: ReactMarkdownOptions) => JSX.Element // to do: remove this. although it breaks our jest tests // Optional callback to save a generated SQL output - onSaveGeneratedSQL?: (answer: string, resolve: any) => void + saveGeneratedSQL?: (answer: string, title: string) => Promise } export const CommandMenuContext = createContext(undefined) export const useCommandMenu = () => { @@ -49,7 +49,7 @@ export interface CommandMenuProviderProps { // to do: remove this prop // this is a temporary hack as ReactMarkdown fails our jest tests if we import the package within this UI package MarkdownHandler: (props: ReactMarkdownOptions) => JSX.Element - onSaveGeneratedSQL?: (answer: string, resolve: any) => void + saveGeneratedSQL?: (answer: string, title: string) => Promise } const CommandMenuProvider = ({ @@ -58,7 +58,7 @@ const CommandMenuProvider = ({ projectRef, apiKeys, MarkdownHandler, - onSaveGeneratedSQL, + saveGeneratedSQL, }: PropsWithChildren) => { const [isOpen, setIsOpen] = useState(false) const [isLoading, setIsLoading] = useState(false) @@ -88,7 +88,7 @@ const CommandMenuProvider = ({ site, project, MarkdownHandler, - onSaveGeneratedSQL, + saveGeneratedSQL, }} > {children} diff --git a/packages/ui/src/components/Command/GenerateSQL.tsx b/packages/ui/src/components/Command/GenerateSQL.tsx index d87821d6f8d..44530d694d0 100644 --- a/packages/ui/src/components/Command/GenerateSQL.tsx +++ b/packages/ui/src/components/Command/GenerateSQL.tsx @@ -10,34 +10,77 @@ import { IconSave, IconUser, Input, + Message, MessageRole, MessageStatus, + queryAi, useAiChat, } from 'ui' +import { stripIndent } from 'common-tags' import CopyToClipboard from 'react-copy-to-clipboard' import { cn } from './../../utils/cn' import { SAMPLE_QUERIES } from './Command.constants' import { AiIcon, AiIconChat } from './Command.icons' import { CommandItem } from './Command.utils' import { useCommandMenu } from './CommandMenuProvider' -import { stripIndent } from 'common-tags' -const SQLOutputActions = ({ answer }: { answer: string }) => { +/** + * Formats a string for use as a title. + * + * Removes punctuation and capitalizes each word + */ +function formatTitle(value: string) { + let words = value.replace(/\.$/, '').replace(/['"]/g, '').split(' ') + words = words.map((word) => { + // Don't capitalize code + if (/[._\(\)]+/.test(word)) { + return word + } + return word.charAt(0).toUpperCase() + word.slice(1) + }) + return words.join(' ') +} + +const SQLOutputActions = ({ answer, messages }: { answer: string; messages: Message[] }) => { const [showCopied, setShowCopied] = useState(false) const [isSaving, setIsSaving] = useState(false) const [isSaved, setIsSaved] = useState(false) - const { project, onSaveGeneratedSQL } = useCommandMenu() - - const applyCallback = () => - onSaveGeneratedSQL !== undefined - ? new Promise((resolve) => onSaveGeneratedSQL(answer, resolve)) - : {} + const { project, saveGeneratedSQL } = useCommandMenu() const onSelectSaveSnippet = async () => { setIsSaving(true) - await applyCallback() + let suggestedTitle + try { + suggestedTitle = await queryAi( + [ + ...messages, + { + role: MessageRole.User, + content: stripIndent` + Generate a title for the above SQL snippet following all of these rules: + - The title is only for the last SQL snippet + - Focus on the main purposes of this snippet + - Use as few words as possible + - Title should be nouns, not verbs + - Do not include word articles (eg. a, the, for, of) + - Do not use words like "SQL" or "snippet" + - Do not output markdown, quotes, etc + - Do not be too verbose + `, + status: MessageStatus.Complete, + }, + ], + 10000 + ) + } catch (err) { + suggestedTitle = '' + } + + const formattedTitle = formatTitle(suggestedTitle) + + await saveGeneratedSQL?.(answer, formattedTitle) setIsSaved(true) setIsSaving(false) } @@ -71,7 +114,7 @@ const SQLOutputActions = ({ answer }: { answer: string }) => { {showCopied ? 'Copied' : 'Copy SQL'} - {project?.ref !== undefined && onSaveGeneratedSQL !== undefined && ( + {project?.ref !== undefined && saveGeneratedSQL !== undefined && (