diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index ad60d7afc75..a17ae6e39eb 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx @@ -1,6 +1,6 @@ import type { UIMessage as MessageType } from '@ai-sdk/react' import { useChat } from '@ai-sdk/react' -import { DefaultChatTransport } from 'ai' +import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai' import { AnimatePresence, motion } from 'framer-motion' import { Eraser, Info, Pencil, X } from 'lucide-react' import { useRouter } from 'next/router' @@ -106,6 +106,21 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { const { mutate: sendEvent } = useSendEventMutation() + const updateMessage = useCallback( + ({ + messageId, + resultId, + results, + }: { + messageId: string + resultId?: string + results: any[] + }) => { + snap.updateMessage({ id: messageId, resultId, results }) + }, + [snap] + ) + // Handle completion of the assistant's response const handleChatFinish = useCallback( ({ message }: { message: MessageType }) => { @@ -119,6 +134,17 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { [snap] ) + // Sanitize messages to remove Valtio proxy wrappers that can't be cloned + const sanitizedMessages = useMemo(() => { + if (!snap.activeChat?.messages) return undefined + + return snap.activeChat.messages.map((msg: any) => { + // Convert proxy objects to plain objects + const plainMessage = JSON.parse(JSON.stringify(msg)) + return plainMessage + }) + }, [snap.activeChat?.messages]) + // TODO(refactor): This useChat hook should be moved down into each chat session. // That way we won't have to disable switching chats while the chat is loading, // and don't run the risk of messages getting mixed up between chats. @@ -133,25 +159,31 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { regenerate, } = useChat({ id: snap.activeChatId, - // [Alaister] typecast is needed here because valtio returns readonly arrays - // and useChat expects a mutable array - messages: snap.activeChat?.messages as unknown as MessageType[] | undefined, + sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls, + messages: sanitizedMessages, async onToolCall({ toolCall }) { + if (toolCall.dynamic) { + return + } + if (toolCall.toolName === 'rename_chat') { const { newName } = toolCall.input as { newName: string } + if (snap.activeChatId && newName?.trim()) { snap.renameChat(snap.activeChatId, newName.trim()) + addToolResult({ tool: toolCall.toolName, toolCallId: toolCall.toolCallId, output: 'Chat renamed', }) + } else { + addToolResult({ + tool: toolCall.toolName, + toolCallId: toolCall.toolCallId, + output: 'Failed to rename chat: Invalid chat or name', + }) } - addToolResult({ - tool: toolCall.toolName, - toolCallId: toolCall.toolCallId, - output: 'Failed to rename chat: Invalid chat or name', - }) } }, transport: new DefaultChatTransport({ @@ -160,7 +192,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { // [Joshen] Specifically limiting the chat history that get's sent to reduce the // size of the context that goes into the model. This should always be an odd number // as much as possible so that the first message is always the user's - const MAX_CHAT_HISTORY = 5 + const MAX_CHAT_HISTORY = 7 const slicedMessages = messages.slice(-MAX_CHAT_HISTORY) @@ -198,21 +230,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming' - const updateMessage = useCallback( - ({ - messageId, - resultId, - results, - }: { - messageId: string - resultId?: string - results: any[] - }) => { - snap.updateMessage({ id: messageId, resultId, results }) - }, - [snap] - ) - const deleteMessageFromHere = useCallback( (messageId: string) => { // Find the message index in current chatMessages @@ -299,7 +316,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { ) const hasMessages = chatMessages.length > 0 - const isShowingOnboarding = !hasMessages && isApiKeySet const sendMessageToAssistant = (finalContent: string) => { if (editingMessageId) { @@ -391,18 +407,17 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { ]} >
- - - - {hasMessages && ( + + {hasMessages ? ( + {renderedMessages} {error && ( @@ -442,14 +457,24 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { )} - )} - - - + + + ) : ( + setValue(val)} + onFocusInput={() => inputRef.current?.focus()} + /> + )} {editingMessageId && ( @@ -519,19 +544,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { /> )} - {isShowingOnboarding && ( - setValue(val)} - onFocusInput={() => inputRef.current?.focus()} - /> - )} - lint.level === LINTER_LEVELS.ERROR) ?? [] + const [securityErrorLints, performanceErrorLints] = partition( + errorLints, + (lint) => lint.categories?.[0] === 'SECURITY' + ) + return ( -
-
-

How can I assist you?

-

- Generate SQL, RLS policies and edge functions, debug issues or check on your project - health. -

-
-
- {prompts.map((item, index) => ( - - - - ))} +
+
+
+

How can I assist you?

+ {suggestions?.prompts?.length ? ( + <> +

Suggestions

+ {prompts.map((item, index) => ( + + + + ))} + + ) : ( + <> + {isLintsLoading ? ( +
+ {Array.from({ length: 6 }).map((_, index) => ( + + ))} +
+ ) : ( + <> + {performanceErrorLints.length > 0 && ( +
+

+ Improve Performance +

+ {performanceErrorLints.map((lint, index) => { + return ( + + ) + })} +
+ )} + + {securityErrorLints.length > 0 && ( +
+

+ Improve Security +

+ {securityErrorLints.map((lint, index) => { + return ( + + ) + })} +
+ )} + +
+

Ideas

+ {prompts.map((item, index) => ( + + ))} +
+ + )} + + )} +
)