import { format } from 'sql-formatter' import { useCallback, useEffect, useRef, useState } from 'react' import { Button, CodeBlock, IconAlertTriangle, IconCornerDownLeft, IconUser, Input, Toggle, MessageRole, MessageStatus, useAiChat, } from 'ui' import { cn } from '../../../utils/cn' import { AiIcon, AiIconChat } from '../Command.icons' import { CommandItem } from '../Command.utils' import { useCommandMenu } from '../CommandMenuProvider' import { SAMPLE_QUERIES } from '../Command.constants' import SQLOutputActions from './SQLOutputActions' import { generatePrompt } from './GenerateSQL.utils' const GenerateSQL = () => { // [Joshen] Temp hack to ensure that generatePrompt receives updated value // of includeSchemaMetadata, needs to be fixed const includeSchemaMetadataRef = useRef() const [includeSchemaMetadata, setIncludeSchemaMetadata] = useState(false) const [selectedCategory, setSelectedCategory] = useState(SAMPLE_QUERIES[0].category) const { isLoading, setIsLoading, search, setSearch, isOptedInToAI, metadata, project } = useCommandMenu() const { flags } = metadata || {} const allowSendingSchemaMetadata = project?.ref !== undefined && flags?.allowCMDKDataOptIn && isOptedInToAI const { submit, reset, messages, isResponding, hasError } = useAiChat({ messageTemplate: (message) => { return generatePrompt( message, isOptedInToAI && includeSchemaMetadataRef.current ? metadata : undefined ) }, setIsLoading, }) const handleSubmit = useCallback( (message: string) => { setSearch('') submit(message) }, [submit] ) const handleReset = useCallback(() => { setSearch('') reset() }, [reset]) useEffect(() => { if (search) handleSubmit(search) includeSchemaMetadataRef.current = includeSchemaMetadata }, []) const formatAnswer = (answer: string) => { try { return format(answer, { language: 'postgresql', keywordCase: 'lower', }) } catch (error: any) { return answer } } return (
e.stopPropagation()}>
{messages.map((message, i) => { switch (message.role) { case MessageRole.User: return (
{message.content}
) case MessageRole.Assistant: const unformattedAnswer = message.content .replace(/```sql/g, '') .replace(/```.*/gs, '') .replace(/-- End of SQL query\.*/g, '') .trim() const answer = message.status === MessageStatus.Complete ? formatAnswer(unformattedAnswer) : unformattedAnswer const cantHelp = answer === "Sorry, I don't know how to help with that." return (
<> {message.status === MessageStatus.Pending ? (
) : cantHelp ? (

Sorry, I don't know how to help with that.

) : (
{answer} {message.status === MessageStatus.Complete && ( )}
)}
) } })} {messages.length === 0 && !hasError && (

Example queries

Use these example queries to help get your project started quickly.

    {SAMPLE_QUERIES.map((item, index) => (
  • setSelectedCategory(item.category)} className={cn( 'px-4 py-1 cursor-pointer text-sm hover:bg-slate-300 rounded-md', selectedCategory === item.category && 'bg-slate-400 ' )} > {item.category}
  • ))}
    {SAMPLE_QUERIES.find((item) => item.category === selectedCategory)?.queries.map( (query, index) => ( { if (!search) { handleSubmit(query) } }} forceMount key={query.replace(/\s+/g, '_')} >

    {query}

    ) )}
)} {hasError && (

Sorry, looks like Clippy is having a hard time!

Please try again in a bit.

)}
{allowSendingSchemaMetadata && (

Include table names, column names and their corresponding data types in prompt

This will generate answers that are more relevant to your project's schema

setIncludeSchemaMetadata((prev) => { includeSchemaMetadataRef.current = !prev return !prev }) } />
)} {!isLoading && !isResponding ? (
Submit message
) : null} } onChange={(e) => { if (!isLoading || !isResponding) { setSearch(e.target.value) } }} onKeyDown={(e) => { switch (e.key) { case 'Enter': if (!search || isLoading || isResponding) return return handleSubmit(search) default: return } }} />
) } export default GenerateSQL