AssistantChatForm.tsx177 lines · main
| 1 | import { useBreakpoint } from 'common' |
| 2 | import { ArrowUp, Loader2, Square } from 'lucide-react' |
| 3 | import { ChangeEvent, FormEvent, forwardRef, KeyboardEvent, memo, useRef } from 'react' |
| 4 | import { ExpandingTextArea } from 'ui' |
| 5 | import { cn } from 'ui/src/lib/utils' |
| 6 | |
| 7 | import { ButtonTooltip } from '../ButtonTooltip' |
| 8 | import { type SqlSnippet } from './AIAssistant.types' |
| 9 | import { ModelSelector } from './ModelSelector' |
| 10 | import { getSnippetContent, SnippetRow } from './SnippetRow' |
| 11 | import type { AssistantModelId } from '@/lib/ai/model.utils' |
| 12 | |
| 13 | export interface FormProps { |
| 14 | /* The ref for the textarea, optional. Exposed for the CommandsPopover to attach events. */ |
| 15 | textAreaRef?: React.RefObject<HTMLTextAreaElement | null> |
| 16 | /* The loading state of the form */ |
| 17 | loading: boolean |
| 18 | /* The disabled state of the form */ |
| 19 | disabled?: boolean |
| 20 | /* The value of the textarea */ |
| 21 | value?: string |
| 22 | /* The function to handle the value change */ |
| 23 | onValueChange: (value: ChangeEvent<HTMLTextAreaElement>) => void |
| 24 | /** |
| 25 | * If true, include SQL snippets in the message sent to onSubmit |
| 26 | */ |
| 27 | includeSnippetsInMessage?: boolean |
| 28 | /** |
| 29 | * The function to handle the form submission |
| 30 | */ |
| 31 | onSubmit: (message: string) => void |
| 32 | /** |
| 33 | * The function to handle stopping the stream |
| 34 | */ |
| 35 | onStop?: () => void |
| 36 | /* The placeholder of the textarea */ |
| 37 | placeholder?: string |
| 38 | /* SQL snippets to display above the form - can be strings or objects with label and content */ |
| 39 | sqlSnippets?: SqlSnippet[] |
| 40 | /* Function to handle removing a SQL snippet */ |
| 41 | onRemoveSnippet?: (index: number) => void |
| 42 | /* Additional class name for the snippets container */ |
| 43 | snippetsClassName?: string |
| 44 | /* Additional class name for the form wrapper */ |
| 45 | className?: string |
| 46 | /* If currently editing an existing message */ |
| 47 | isEditing?: boolean |
| 48 | /* The currently selected AI model */ |
| 49 | selectedModel: AssistantModelId |
| 50 | /* Callback when a model is chosen */ |
| 51 | onSelectModel: (model: AssistantModelId) => void |
| 52 | } |
| 53 | |
| 54 | const AssistantChatFormComponent = forwardRef<HTMLFormElement, FormProps>( |
| 55 | ( |
| 56 | { |
| 57 | loading = false, |
| 58 | disabled = false, |
| 59 | value = '', |
| 60 | textAreaRef, |
| 61 | onValueChange, |
| 62 | onSubmit, |
| 63 | onStop, |
| 64 | placeholder, |
| 65 | sqlSnippets, |
| 66 | onRemoveSnippet, |
| 67 | snippetsClassName, |
| 68 | includeSnippetsInMessage = false, |
| 69 | className, |
| 70 | isEditing = false, |
| 71 | selectedModel, |
| 72 | onSelectModel, |
| 73 | ...props |
| 74 | }, |
| 75 | _ref |
| 76 | ) => { |
| 77 | const formRef = useRef<HTMLFormElement>(null) |
| 78 | const isMobile = useBreakpoint('md') |
| 79 | |
| 80 | const handleSubmit = (event?: FormEvent<HTMLFormElement>) => { |
| 81 | if (event) event.preventDefault() |
| 82 | if (disabled || !value || (loading && !isEditing)) return |
| 83 | |
| 84 | let finalMessage = value |
| 85 | if (includeSnippetsInMessage && sqlSnippets && sqlSnippets.length > 0) { |
| 86 | const sqlSnippetsString = sqlSnippets |
| 87 | .map((snippet: SqlSnippet) => '```sql\n' + getSnippetContent(snippet) + '\n```') |
| 88 | .join('\n') |
| 89 | finalMessage = [value, sqlSnippetsString].filter(Boolean).join('\n\n') |
| 90 | } |
| 91 | |
| 92 | onSubmit(finalMessage) |
| 93 | } |
| 94 | |
| 95 | const handleKeyDown = (event: KeyboardEvent<HTMLTextAreaElement>) => { |
| 96 | if (event.key === 'Enter' && !event.shiftKey) { |
| 97 | event.preventDefault() |
| 98 | handleSubmit() |
| 99 | } |
| 100 | } |
| 101 | |
| 102 | const canSubmit = !disabled && !loading && !!value |
| 103 | |
| 104 | return ( |
| 105 | <div className="w-full"> |
| 106 | <form |
| 107 | id="assistant-chat" |
| 108 | ref={formRef} |
| 109 | {...props} |
| 110 | onSubmit={handleSubmit} |
| 111 | className={cn('relative', className)} |
| 112 | > |
| 113 | {sqlSnippets && sqlSnippets.length > 0 && ( |
| 114 | <SnippetRow |
| 115 | snippets={sqlSnippets} |
| 116 | onRemoveSnippet={onRemoveSnippet} |
| 117 | className="absolute top-0 left-0 right-0 px-1.5 py-1.5" |
| 118 | /> |
| 119 | )} |
| 120 | <ExpandingTextArea |
| 121 | autoFocus={!isMobile} |
| 122 | ref={textAreaRef} |
| 123 | disabled={disabled} |
| 124 | className={cn( |
| 125 | 'text-base md:text-sm pr-10 pb-9 max-h-64', |
| 126 | sqlSnippets && sqlSnippets.length > 0 && 'pt-10' |
| 127 | )} |
| 128 | placeholder={placeholder} |
| 129 | spellCheck={false} |
| 130 | rows={3} |
| 131 | value={value} |
| 132 | onChange={(event) => onValueChange(event)} |
| 133 | onKeyDown={handleKeyDown} |
| 134 | /> |
| 135 | <div className="absolute inset-x-1.5 bottom-1.5 flex items-center justify-between pointer-events-none"> |
| 136 | <div className="pointer-events-auto"> |
| 137 | <ModelSelector selectedModel={selectedModel} onSelectModel={onSelectModel} /> |
| 138 | </div> |
| 139 | |
| 140 | <div className="flex gap-3 items-center pointer-events-auto"> |
| 141 | {loading ? ( |
| 142 | onStop ? ( |
| 143 | <ButtonTooltip |
| 144 | type="outline" |
| 145 | aria-label="Stop response" |
| 146 | icon={<Square fill="currentColor" className="scale-75" />} |
| 147 | onClick={onStop} |
| 148 | className="w-7 h-7 rounded-full p-0 text-center flex items-center justify-center" |
| 149 | tooltip={{ content: { side: 'top', text: 'Stop response' } }} |
| 150 | /> |
| 151 | ) : ( |
| 152 | <Loader2 size={22} className="animate-spin size-7 text-muted" strokeWidth={1} /> |
| 153 | ) |
| 154 | ) : ( |
| 155 | <ButtonTooltip |
| 156 | htmlType="submit" |
| 157 | aria-label="Send message" |
| 158 | icon={<ArrowUp />} |
| 159 | disabled={!canSubmit} |
| 160 | className={cn( |
| 161 | 'w-7 h-7 rounded-full p-0 text-center flex items-center justify-center', |
| 162 | !canSubmit ? 'opacity-50' : 'opacity-100' |
| 163 | )} |
| 164 | tooltip={{ content: { side: 'top', text: 'Send message' } }} |
| 165 | /> |
| 166 | )} |
| 167 | </div> |
| 168 | </div> |
| 169 | </form> |
| 170 | </div> |
| 171 | ) |
| 172 | } |
| 173 | ) |
| 174 | |
| 175 | AssistantChatFormComponent.displayName = 'AssistantChatFormComponent' |
| 176 | |
| 177 | export const AssistantChatForm = memo(AssistantChatFormComponent) |