AssistantChatForm.tsx177 lines · main
1import { useBreakpoint } from 'common'
2import { ArrowUp, Loader2, Square } from 'lucide-react'
3import { ChangeEvent, FormEvent, forwardRef, KeyboardEvent, memo, useRef } from 'react'
4import { ExpandingTextArea } from 'ui'
5import { cn } from 'ui/src/lib/utils'
6
7import { ButtonTooltip } from '../ButtonTooltip'
8import { type SqlSnippet } from './AIAssistant.types'
9import { ModelSelector } from './ModelSelector'
10import { getSnippetContent, SnippetRow } from './SnippetRow'
11import type { AssistantModelId } from '@/lib/ai/model.utils'
12
13export 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
54const 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
175AssistantChatFormComponent.displayName = 'AssistantChatFormComponent'
176
177export const AssistantChatForm = memo(AssistantChatFormComponent)