Skip to content

Commit 06b4210

Browse files
gregnrgregnr
authored
Merge pull request #138 from supabase-community/feat/byo-llm
feat: bring your own llm
2 parents 3e31c4e + 66dac4f commit 06b4210

27 files changed

+1837
-825
lines changed

‎apps/postgres-new/.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
public/sw.mjs

‎apps/postgres-new/app/api/chat/route.ts

Lines changed: 5 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@ import { createOpenAI } from '@ai-sdk/openai'
22
import { Ratelimit } from '@upstash/ratelimit'
33
import { kv } from '@vercel/kv'
44
import { convertToCoreMessages, streamText, ToolInvocation, ToolResultPart } from 'ai'
5-
import { codeBlock } from 'common-tags'
6-
import { convertToCoreTools, maxMessageContext, maxRowLimit, tools } from '~/lib/tools'
5+
import { getSystemPrompt } from '~/lib/system-prompt'
6+
import { convertToCoreTools, maxMessageContext, tools } from '~/lib/tools'
77
import { createClient } from '~/utils/supabase/server'
88
import { ChatInferenceEventToolResult, logEvent } from '~/utils/telemetry'
99

@@ -72,49 +72,8 @@ export async function POST(req: Request) {
7272
const coreMessages = convertToCoreMessages(trimmedMessageContext)
7373
const coreTools = convertToCoreTools(tools)
7474

75-
const result = await streamText({
76-
system: codeBlock`
77-
You are a helpful database assistant. Under the hood you have access to an in-browser Postgres database called PGlite (https://github.com/electric-sql/pglite).
78-
Some special notes about this database:
79-
- foreign data wrappers are not supported
80-
- the following extensions are available:
81-
- plpgsql [pre-enabled]
82-
- vector (https://github.com/pgvector/pgvector) [pre-enabled]
83-
- use <=> for cosine distance (default to this)
84-
- use <#> for negative inner product
85-
- use <-> for L2 distance
86-
- use <+> for L1 distance
87-
- note queried vectors will be truncated/redacted due to their size - export as CSV if the full vector is required
88-
89-
When generating tables, do the following:
90-
- For primary keys, always use "id bigint primary key generated always as identity" (not serial)
91-
- Prefer 'text' over 'varchar'
92-
- Keep explanations brief but helpful
93-
- Don't repeat yourself after creating the table
94-
95-
When creating sample data:
96-
- Make the data realistic, including joined data
97-
- Check for existing records/conflicts in the table
98-
99-
When querying data, limit to 5 by default. The maximum number of rows you're allowed to fetch is ${maxRowLimit} (to protect AI from token abuse).
100-
If the user needs to fetch more than ${maxRowLimit} rows at once, they can export the query as a CSV.
101-
102-
When performing FTS, always use 'simple' (languages aren't available).
103-
104-
When importing CSVs try to solve the problem yourself (eg. use a generic text column, then refine)
105-
vs. asking the user to change the CSV. No need to select rows after importing.
106-
107-
You also know math. All math equations and expressions must be written in KaTex and must be wrapped in double dollar \`$$\`:
108-
- Inline: $$\\sqrt{26}$$
109-
- Multiline:
110-
$$
111-
\\sqrt{26}
112-
$$
113-
114-
No images are allowed. Do not try to generate or link images, including base64 data URLs.
115-
116-
Feel free to suggest corrections for suspected typos.
117-
`,
75+
const result = streamText({
76+
system: getSystemPrompt(),
11877
model: openai(chatModel),
11978
messages: coreMessages,
12079
tools: coreTools,
@@ -158,7 +117,7 @@ export async function POST(req: Request) {
158117
},
159118
})
160119

161-
return result.toAIStreamResponse()
120+
return result.toDataStreamResponse()
162121
}
163122

164123
function getEventToolResult(toolResult: ToolResultPart): ChatInferenceEventToolResult | undefined {

‎apps/postgres-new/components/app-provider.tsx

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import {
3232
import { legacyDomainHostname } from '~/lib/util'
3333
import { parse, serialize } from '~/lib/websocket-protocol'
3434
import { createClient } from '~/utils/supabase/client'
35+
import { useModelProvider } from './model-provider/use-model-provider'
3536

3637
export type AppProps = PropsWithChildren
3738

@@ -252,6 +253,9 @@ export default function AppProvider({ children }: AppProps) {
252253
const [isLegacyDomain, setIsLegacyDomain] = useState(false)
253254
const [isLegacyDomainRedirect, setIsLegacyDomainRedirect] = useState(false)
254255

256+
const [modelProviderError, setModelProviderError] = useState<string>()
257+
const [isModelProviderDialogOpen, setIsModelProviderDialogOpen] = useState(false)
258+
255259
useEffect(() => {
256260
const isLegacyDomain = window.location.hostname === legacyDomainHostname
257261
const urlParams = new URLSearchParams(window.location.search)
@@ -263,12 +267,17 @@ export default function AppProvider({ children }: AppProps) {
263267
setIsRenameDialogOpen(isLegacyDomain || isLegacyDomainRedirect)
264268
}, [])
265269

270+
const modelProvider = useModelProvider()
271+
266272
return (
267273
<AppContext.Provider
268274
value={{
269275
user,
270276
isLoadingUser,
271277
liveShare,
278+
modelProvider,
279+
modelProviderError,
280+
setModelProviderError,
272281
signIn,
273282
signOut,
274283
isSignInDialogOpen,
@@ -277,6 +286,8 @@ export default function AppProvider({ children }: AppProps) {
277286
setIsRenameDialogOpen,
278287
isRateLimited,
279288
setIsRateLimited,
289+
isModelProviderDialogOpen,
290+
setIsModelProviderDialogOpen,
280291
focusRef,
281292
dbManager,
282293
pgliteVersion,
@@ -305,6 +316,8 @@ export type AppContextValues = {
305316
setIsRenameDialogOpen: (open: boolean) => void
306317
isRateLimited: boolean
307318
setIsRateLimited: (limited: boolean) => void
319+
isModelProviderDialogOpen: boolean
320+
setIsModelProviderDialogOpen: (open: boolean) => void
308321
focusRef: RefObject<FocusHandle>
309322
dbManager?: DbManager
310323
pgliteVersion?: string
@@ -316,6 +329,9 @@ export type AppContextValues = {
316329
clientIp: string | null
317330
isLiveSharing: boolean
318331
}
332+
modelProvider: ReturnType<typeof useModelProvider>
333+
modelProviderError?: string
334+
setModelProviderError: (error: string | undefined) => void
319335
isLegacyDomain: boolean
320336
isLegacyDomainRedirect: boolean
321337
}

‎apps/postgres-new/components/byo-llm-button.tsx

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import { Brain } from 'lucide-react'
2+
import { useApp } from '~/components/app-provider'
3+
import { Button } from '~/components/ui/button'
4+
5+
export type ByoLlmButtonProps = {
6+
onClick?: () => void
7+
}
8+
9+
export default function ByoLlmButton({ onClick }: ByoLlmButtonProps) {
10+
const { setIsModelProviderDialogOpen } = useApp()
11+
12+
return (
13+
<Button
14+
className="gap-2 text-base"
15+
onClick={() => {
16+
onClick?.()
17+
setIsModelProviderDialogOpen(true)
18+
}}
19+
>
20+
<Brain size={18} strokeWidth={2} />
21+
Bring your own LLM
22+
</Button>
23+
)
24+
}

‎apps/postgres-new/components/chat.tsx

Lines changed: 61 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import { Message, generateId } from 'ai'
44
import { useChat } from 'ai/react'
55
import { AnimatePresence, m } from 'framer-motion'
6-
import { ArrowDown, ArrowUp, Flame, Paperclip, PlugIcon, Square } from 'lucide-react'
6+
import { AlertCircle, ArrowDown, ArrowUp, Flame, Paperclip, PlugIcon, Square } from 'lucide-react'
77
import {
88
FormEventHandler,
99
useCallback,
@@ -22,6 +22,7 @@ import { requestFileUpload } from '~/lib/util'
2222
import { cn } from '~/lib/utils'
2323
import { AiIconAnimation } from './ai-icon-animation'
2424
import { useApp } from './app-provider'
25+
import ByoLlmButton from './byo-llm-button'
2526
import ChatMessage from './chat-message'
2627
import { CopyableField } from './copyable-field'
2728
import SignInButton from './sign-in-button'
@@ -51,8 +52,17 @@ export function getInitialMessages(tables: TablesData): Message[] {
5152
}
5253

5354
export default function Chat() {
54-
const { user, isLoadingUser, focusRef, setIsSignInDialogOpen, isRateLimited, liveShare } =
55-
useApp()
55+
const {
56+
user,
57+
isLoadingUser,
58+
focusRef,
59+
setIsSignInDialogOpen,
60+
isRateLimited,
61+
liveShare,
62+
modelProvider,
63+
modelProviderError,
64+
setIsModelProviderDialogOpen,
65+
} = useApp()
5666
const [inputFocusState, setInputFocusState] = useState(false)
5767

5868
const {
@@ -155,7 +165,7 @@ export default function Chat() {
155165
cursor: dropZoneCursor,
156166
} = useDropZone({
157167
async onDrop(files) {
158-
if (!user) {
168+
if (isAuthRequired) {
159169
return
160170
}
161171

@@ -223,8 +233,10 @@ export default function Chat() {
223233

224234
const [isMessageAnimationComplete, setIsMessageAnimationComplete] = useState(false)
225235

236+
const isAuthRequired = user === undefined && modelProvider.state?.enabled !== true
237+
226238
const isChatEnabled =
227-
!isLoadingMessages && !isLoadingSchema && user !== undefined && !liveShare.isLiveSharing
239+
!isLoadingMessages && !isLoadingSchema && !isAuthRequired && !liveShare.isLiveSharing
228240

229241
const isSubmitEnabled = isChatEnabled && Boolean(input.trim())
230242

@@ -293,6 +305,42 @@ export default function Chat() {
293305
isLast={i === messages.length - 1}
294306
/>
295307
))}
308+
<AnimatePresence initial={false}>
309+
{modelProviderError && !isLoading && (
310+
<m.div
311+
layout="position"
312+
className="flex flex-col gap-4 justify-start items-center max-w-96 p-4 bg-destructive rounded-md text-sm"
313+
variants={{
314+
hidden: { scale: 0 },
315+
show: { scale: 1, transition: { delay: 0.5 } },
316+
}}
317+
initial="hidden"
318+
animate="show"
319+
exit="hidden"
320+
>
321+
<AlertCircle size={64} strokeWidth={1} />
322+
<div className="flex flex-col items-center text-start gap-4">
323+
<h3 className="font-bold">Whoops!</h3>
324+
<p className="text-center">
325+
There was an error connecting to your custom model provider:{' '}
326+
{modelProviderError}
327+
</p>
328+
<p>
329+
Double check your{' '}
330+
<a
331+
className="underline cursor-pointer"
332+
onClick={() => {
333+
setIsModelProviderDialogOpen(true)
334+
}}
335+
>
336+
API info
337+
</a>
338+
.
339+
</p>
340+
</div>
341+
</m.div>
342+
)}
343+
</AnimatePresence>
296344
<AnimatePresence initial={false}>
297345
{isRateLimited && !isLoading && (
298346
<m.div
@@ -357,7 +405,7 @@ export default function Chat() {
357405
</div>
358406
) : (
359407
<div className="h-full w-full max-w-4xl flex flex-col gap-10 justify-center items-center">
360-
{user ? (
408+
{!isAuthRequired ? (
361409
<>
362410
<LiveShareOverlay databaseId={databaseId} />
363411
<m.h3
@@ -384,11 +432,10 @@ export default function Chat() {
384432
animate="show"
385433
>
386434
<SignInButton />
387-
<p className="font-lighter text-center">
388-
To prevent abuse we ask you to sign in before chatting with AI.
389-
</p>
435+
or
436+
<ByoLlmButton />
390437
<p
391-
className="underline cursor-pointer text-primary/50"
438+
className="underline cursor-pointer text-sm text-primary/50"
392439
onClick={() => {
393440
setIsSignInDialogOpen(true)
394441
}}
@@ -427,7 +474,7 @@ export default function Chat() {
427474
</div>
428475
<div className="flex flex-col items-center gap-3 pb-1 relative">
429476
<AnimatePresence>
430-
{!user && !isLoadingUser && isConversationStarted && (
477+
{isAuthRequired && !isLoadingUser && isConversationStarted && (
431478
<m.div
432479
className="flex flex-col items-center gap-4 max-w-lg my-4"
433480
variants={{
@@ -438,9 +485,8 @@ export default function Chat() {
438485
exit="hidden"
439486
>
440487
<SignInButton />
441-
<p className="font-lighter text-center text-sm">
442-
To prevent abuse we ask you to sign in before chatting with AI.
443-
</p>
488+
or
489+
<ByoLlmButton />
444490
<p
445491
className="underline cursor-pointer text-sm text-primary/50"
446492
onClick={() => {
@@ -487,7 +533,7 @@ export default function Chat() {
487533
onClick={async (e) => {
488534
e.preventDefault()
489535

490-
if (!user) {
536+
if (isAuthRequired) {
491537
return
492538
}
493539

0 commit comments

Comments
 (0)