Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 68 additions & 56 deletions apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
Original file line number Diff line number Diff line change
@@ -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'
Expand Down Expand Up @@ -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 }) => {
Expand All @@ -122,6 +137,17 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
// 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.
// 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])

const {
messages: chatMessages,
status: chatStatus,
Expand All @@ -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({
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -299,7 +316,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
)

const hasMessages = chatMessages.length > 0
const isShowingOnboarding = !hasMessages && isApiKeySet

const sendMessageToAssistant = (finalContent: string) => {
if (editingMessageId) {
Expand Down Expand Up @@ -391,18 +407,17 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
]}
>
<div className={cn('flex flex-col h-full', className)}>
<Conversation className={cn('flex-1')}>
<AIAssistantHeader
isChatLoading={isChatLoading}
onNewChat={snap.newChat}
onCloseAssistant={snap.closeAssistant}
showMetadataWarning={showMetadataWarning}
updatedOptInSinceMCP={updatedOptInSinceMCP}
isHipaaProjectDisallowed={isHipaaProjectDisallowed as boolean}
aiOptInLevel={aiOptInLevel}
/>

{hasMessages && (
<AIAssistantHeader
isChatLoading={isChatLoading}
onNewChat={snap.newChat}
onCloseAssistant={snap.closeAssistant}
showMetadataWarning={showMetadataWarning}
updatedOptInSinceMCP={updatedOptInSinceMCP}
isHipaaProjectDisallowed={isHipaaProjectDisallowed as boolean}
aiOptInLevel={aiOptInLevel}
/>
{hasMessages ? (
<Conversation className={cn('flex-1')}>
<ConversationContent className="w-full px-7 py-8 mb-10">
{renderedMessages}
{error && (
Expand Down Expand Up @@ -442,14 +457,24 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
<motion.span
animate={{ opacity: [1, 0] }}
transition={{ duration: 1, repeat: Infinity, ease: 'linear' }}
className="inline-block w-1.5 h-4 bg-foreground-lighter"
className="inline-block w-1.5 h-4 bg-foreground-lighter mt-4"
/>
)}
</ConversationContent>
)}

<ConversationScrollButton />
</Conversation>
<ConversationScrollButton />
</Conversation>
) : (
<AIOnboarding
sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}
suggestions={
snap.suggestions as
| { title?: string; prompts?: { label: string; description: string }[] }
| undefined
}
onValueChange={(val) => setValue(val)}
onFocusInput={() => inputRef.current?.focus()}
/>
)}

<AnimatePresence>
{editingMessageId && (
Expand Down Expand Up @@ -519,19 +544,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
/>
)}

{isShowingOnboarding && (
<AIOnboarding
sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}
suggestions={
snap.suggestions as
| { title?: string; prompts?: { label: string; description: string }[] }
| undefined
}
onValueChange={(val) => setValue(val)}
onFocusInput={() => inputRef.current?.focus()}
/>
)}

<AssistantChatForm
textAreaRef={inputRef}
className={cn(
Expand Down
172 changes: 140 additions & 32 deletions apps/studio/components/ui/AIAssistantPanel/AIOnboarding.tsx
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
import { motion } from 'framer-motion'
import { FileText } from 'lucide-react'
import { partition } from 'lodash'
import { BarChart, FileText, Shield } from 'lucide-react'

import { Button } from 'ui'
import { Button, Skeleton } from 'ui'
import { useParams } from 'common'
import { LINTER_LEVELS } from 'components/interfaces/Linter/Linter.constants'
import { createLintSummaryPrompt } from 'components/interfaces/Linter/Linter.utils'
import { useProjectLintsQuery } from 'data/lint/lint-query'
import { type SqlSnippet } from './AIAssistant.types'
import { codeSnippetPrompts, defaultPrompts } from './AIAssistant.prompts'

Expand Down Expand Up @@ -31,37 +36,140 @@ export const AIOnboarding = ({
? codeSnippetPrompts
: defaultPrompts

const { ref: projectRef } = useParams()
const {
data: lints,
isLoading: isLoadingLints,
isFetching: isFetchingLints,
} = useProjectLintsQuery({ projectRef })
const isLintsLoading = isLoadingLints || isFetchingLints

const errorLints = lints?.filter((lint) => lint.level === LINTER_LEVELS.ERROR) ?? []
const [securityErrorLints, performanceErrorLints] = partition(
errorLints,
(lint) => lint.categories?.[0] === 'SECURITY'
)

return (
<div className="w-full mb-6">
<div className="px-4 mb-4">
<h2 className="heading-section text-foreground mb-1">How can I assist you?</h2>
<p className="text-foreground-light text-sm">
Generate SQL, RLS policies and edge functions, debug issues or check on your project
health.
</p>
</div>
<div>
{prompts.map((item, index) => (
<motion.div
key={index}
initial={{ y: 5, opacity: 0 }}
animate={{ y: 0, opacity: 1 }}
transition={{ duration: 0.3, delay: index * 0.05 }}
>
<Button
size="small"
type="text"
className="w-full justify-start"
icon={<FileText strokeWidth={1.5} size={14} className="text-foreground-light" />}
onClick={() => {
onValueChange(item.prompt)
onFocusInput?.()
}}
>
{item.title}
</Button>
</motion.div>
))}
<div className="flex-1 overflow-y-auto">
<div className="w-full flex-1 max-h-full min-h-full px-4 flex flex-col gap-0">
<div className="mt-auto w-full space-y-6 py-8 ">
<h2 className="heading-section text-foreground mx-4">How can I assist you?</h2>
{suggestions?.prompts?.length ? (
<>
<h3 className="heading-meta text-foreground-light mb-3 mx-4">Suggestions</h3>
{prompts.map((item, index) => (
<motion.div
key={index}
initial={{ y: 5, opacity: 0 }}
animate={{ y: 0, opacity: 1 }}
transition={{ duration: 0.3, delay: index * 0.05 }}
>
<Button
size="small"
type="text"
className="w-full justify-start border-b hover:border-b-0 hover:rounded-md rounded-none"
icon={
<FileText strokeWidth={1.5} size={14} className="text-foreground-light" />
}
onClick={() => {
onValueChange(item.prompt)
onFocusInput?.()
}}
>
{item.title}
</Button>
</motion.div>
))}
</>
) : (
<>
{isLintsLoading ? (
<div className="px-4 flex flex-col gap-2">
{Array.from({ length: 6 }).map((_, index) => (
<Skeleton className="h-4 w-full" />
))}
</div>
) : (
<>
{performanceErrorLints.length > 0 && (
<div className="mb-4">
<h3 className="heading-meta text-foreground-light mb-3 mx-4">
Improve Performance
</h3>
{performanceErrorLints.map((lint, index) => {
return (
<Button
size="small"
type="text"
className="w-full justify-start"
icon={
<BarChart
strokeWidth={1.5}
size={14}
className="text-foreground-light"
/>
}
onClick={() => {
onValueChange(createLintSummaryPrompt(lint))
onFocusInput?.()
}}
>
{lint.detail ? lint.detail.replace('\\`', '') : lint.title}
</Button>
)
})}
</div>
)}

{securityErrorLints.length > 0 && (
<div className="mb-4">
<h3 className="heading-meta text-foreground-light mb-3 mx-4">
Improve Security
</h3>
{securityErrorLints.map((lint, index) => {
return (
<Button
size="small"
type="text"
className="w-full justify-start"
icon={<Shield strokeWidth={1.5} size={14} className="text-warning" />}
onClick={() => {
onValueChange(createLintSummaryPrompt(lint))
onFocusInput?.()
}}
>
{lint.detail ? lint.detail.replace(/\\`/g, '') : lint.title}
</Button>
)
})}
</div>
)}

<div>
<h3 className="heading-meta text-foreground-light mb-3 mx-4">Ideas</h3>
{prompts.map((item, index) => (
<Button
size="small"
type="text"
className="w-full justify-start"
icon={
<FileText strokeWidth={1.5} size={14} className="text-foreground-light" />
}
onClick={() => {
onValueChange(item.prompt)
onFocusInput?.()
}}
>
{item.title}
</Button>
))}
</div>
</>
)}
</>
)}
</div>
</div>
</div>
)
Expand Down
Loading
Loading