diff --git a/apps/postgres-new/.gitignore b/apps/postgres-new/.gitignore new file mode 100644 index 00000000..33a7894a --- /dev/null +++ b/apps/postgres-new/.gitignore @@ -0,0 +1 @@ +public/sw.mjs \ No newline at end of file diff --git a/apps/postgres-new/app/api/chat/route.ts b/apps/postgres-new/app/api/chat/route.ts index 4a7a0b62..1488ef0f 100644 --- a/apps/postgres-new/app/api/chat/route.ts +++ b/apps/postgres-new/app/api/chat/route.ts @@ -2,8 +2,8 @@ import { createOpenAI } from '@ai-sdk/openai' import { Ratelimit } from '@upstash/ratelimit' import { kv } from '@vercel/kv' import { convertToCoreMessages, streamText, ToolInvocation, ToolResultPart } from 'ai' -import { codeBlock } from 'common-tags' -import { convertToCoreTools, maxMessageContext, maxRowLimit, tools } from '~/lib/tools' +import { getSystemPrompt } from '~/lib/system-prompt' +import { convertToCoreTools, maxMessageContext, tools } from '~/lib/tools' import { createClient } from '~/utils/supabase/server' import { ChatInferenceEventToolResult, logEvent } from '~/utils/telemetry' @@ -72,49 +72,8 @@ export async function POST(req: Request) { const coreMessages = convertToCoreMessages(trimmedMessageContext) const coreTools = convertToCoreTools(tools) - const result = await streamText({ - system: codeBlock` - 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). - Some special notes about this database: - - foreign data wrappers are not supported - - the following extensions are available: - - plpgsql [pre-enabled] - - vector (https://github.com/pgvector/pgvector) [pre-enabled] - - use <=> for cosine distance (default to this) - - use <#> for negative inner product - - use <-> for L2 distance - - use <+> for L1 distance - - note queried vectors will be truncated/redacted due to their size - export as CSV if the full vector is required - - When generating tables, do the following: - - For primary keys, always use "id bigint primary key generated always as identity" (not serial) - - Prefer 'text' over 'varchar' - - Keep explanations brief but helpful - - Don't repeat yourself after creating the table - - When creating sample data: - - Make the data realistic, including joined data - - Check for existing records/conflicts in the table - - 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). - If the user needs to fetch more than ${maxRowLimit} rows at once, they can export the query as a CSV. - - When performing FTS, always use 'simple' (languages aren't available). - - When importing CSVs try to solve the problem yourself (eg. use a generic text column, then refine) - vs. asking the user to change the CSV. No need to select rows after importing. - - You also know math. All math equations and expressions must be written in KaTex and must be wrapped in double dollar \`$$\`: - - Inline: $$\\sqrt{26}$$ - - Multiline: - $$ - \\sqrt{26} - $$ - - No images are allowed. Do not try to generate or link images, including base64 data URLs. - - Feel free to suggest corrections for suspected typos. - `, + const result = streamText({ + system: getSystemPrompt(), model: openai(chatModel), messages: coreMessages, tools: coreTools, @@ -158,7 +117,7 @@ export async function POST(req: Request) { }, }) - return result.toAIStreamResponse() + return result.toDataStreamResponse() } function getEventToolResult(toolResult: ToolResultPart): ChatInferenceEventToolResult | undefined { diff --git a/apps/postgres-new/components/app-provider.tsx b/apps/postgres-new/components/app-provider.tsx index b4905ced..557a3508 100644 --- a/apps/postgres-new/components/app-provider.tsx +++ b/apps/postgres-new/components/app-provider.tsx @@ -32,6 +32,7 @@ import { import { legacyDomainHostname } from '~/lib/util' import { parse, serialize } from '~/lib/websocket-protocol' import { createClient } from '~/utils/supabase/client' +import { useModelProvider } from './model-provider/use-model-provider' export type AppProps = PropsWithChildren @@ -252,6 +253,9 @@ export default function AppProvider({ children }: AppProps) { const [isLegacyDomain, setIsLegacyDomain] = useState(false) const [isLegacyDomainRedirect, setIsLegacyDomainRedirect] = useState(false) + const [modelProviderError, setModelProviderError] = useState() + const [isModelProviderDialogOpen, setIsModelProviderDialogOpen] = useState(false) + useEffect(() => { const isLegacyDomain = window.location.hostname === legacyDomainHostname const urlParams = new URLSearchParams(window.location.search) @@ -263,12 +267,17 @@ export default function AppProvider({ children }: AppProps) { setIsRenameDialogOpen(isLegacyDomain || isLegacyDomainRedirect) }, []) + const modelProvider = useModelProvider() + return ( void isRateLimited: boolean setIsRateLimited: (limited: boolean) => void + isModelProviderDialogOpen: boolean + setIsModelProviderDialogOpen: (open: boolean) => void focusRef: RefObject dbManager?: DbManager pgliteVersion?: string @@ -316,6 +329,9 @@ export type AppContextValues = { clientIp: string | null isLiveSharing: boolean } + modelProvider: ReturnType + modelProviderError?: string + setModelProviderError: (error: string | undefined) => void isLegacyDomain: boolean isLegacyDomainRedirect: boolean } diff --git a/apps/postgres-new/components/byo-llm-button.tsx b/apps/postgres-new/components/byo-llm-button.tsx new file mode 100644 index 00000000..d09e2072 --- /dev/null +++ b/apps/postgres-new/components/byo-llm-button.tsx @@ -0,0 +1,24 @@ +import { Brain } from 'lucide-react' +import { useApp } from '~/components/app-provider' +import { Button } from '~/components/ui/button' + +export type ByoLlmButtonProps = { + onClick?: () => void +} + +export default function ByoLlmButton({ onClick }: ByoLlmButtonProps) { + const { setIsModelProviderDialogOpen } = useApp() + + return ( + + ) +} diff --git a/apps/postgres-new/components/chat.tsx b/apps/postgres-new/components/chat.tsx index 19cb8750..05a39690 100644 --- a/apps/postgres-new/components/chat.tsx +++ b/apps/postgres-new/components/chat.tsx @@ -3,7 +3,7 @@ import { Message, generateId } from 'ai' import { useChat } from 'ai/react' import { AnimatePresence, m } from 'framer-motion' -import { ArrowDown, ArrowUp, Flame, Paperclip, PlugIcon, Square } from 'lucide-react' +import { AlertCircle, ArrowDown, ArrowUp, Flame, Paperclip, PlugIcon, Square } from 'lucide-react' import { FormEventHandler, useCallback, @@ -22,6 +22,7 @@ import { requestFileUpload } from '~/lib/util' import { cn } from '~/lib/utils' import { AiIconAnimation } from './ai-icon-animation' import { useApp } from './app-provider' +import ByoLlmButton from './byo-llm-button' import ChatMessage from './chat-message' import { CopyableField } from './copyable-field' import SignInButton from './sign-in-button' @@ -51,8 +52,17 @@ export function getInitialMessages(tables: TablesData): Message[] { } export default function Chat() { - const { user, isLoadingUser, focusRef, setIsSignInDialogOpen, isRateLimited, liveShare } = - useApp() + const { + user, + isLoadingUser, + focusRef, + setIsSignInDialogOpen, + isRateLimited, + liveShare, + modelProvider, + modelProviderError, + setIsModelProviderDialogOpen, + } = useApp() const [inputFocusState, setInputFocusState] = useState(false) const { @@ -155,7 +165,7 @@ export default function Chat() { cursor: dropZoneCursor, } = useDropZone({ async onDrop(files) { - if (!user) { + if (isAuthRequired) { return } @@ -223,8 +233,10 @@ export default function Chat() { const [isMessageAnimationComplete, setIsMessageAnimationComplete] = useState(false) + const isAuthRequired = user === undefined && modelProvider.state?.enabled !== true + const isChatEnabled = - !isLoadingMessages && !isLoadingSchema && user !== undefined && !liveShare.isLiveSharing + !isLoadingMessages && !isLoadingSchema && !isAuthRequired && !liveShare.isLiveSharing const isSubmitEnabled = isChatEnabled && Boolean(input.trim()) @@ -293,6 +305,42 @@ export default function Chat() { isLast={i === messages.length - 1} /> ))} + + {modelProviderError && !isLoading && ( + + +
+

Whoops!

+

+ There was an error connecting to your custom model provider:{' '} + {modelProviderError} +

+

+ Double check your{' '} + { + setIsModelProviderDialogOpen(true) + }} + > + API info + + . +

+
+ + )} + {isRateLimited && !isLoading && ( ) : (
- {user ? ( + {!isAuthRequired ? ( <> -

- To prevent abuse we ask you to sign in before chatting with AI. -

+ or +

{ setIsSignInDialogOpen(true) }} @@ -427,7 +474,7 @@ export default function Chat() {

- {!user && !isLoadingUser && isConversationStarted && ( + {isAuthRequired && !isLoadingUser && isConversationStarted && ( -

- To prevent abuse we ask you to sign in before chatting with AI. -

+ or +

{ @@ -487,7 +533,7 @@ export default function Chat() { onClick={async (e) => { e.preventDefault() - if (!user) { + if (isAuthRequired) { return } diff --git a/apps/postgres-new/components/layout.tsx b/apps/postgres-new/components/layout.tsx index 11000a04..06283d37 100644 --- a/apps/postgres-new/components/layout.tsx +++ b/apps/postgres-new/components/layout.tsx @@ -3,14 +3,12 @@ import 'chart.js/auto' import 'chartjs-adapter-date-fns' -import { DialogTrigger } from '@radix-ui/react-dialog' import { LazyMotion, m } from 'framer-motion' -import { Loader, MoreVertical } from 'lucide-react' +import { Loader } from 'lucide-react' import Link from 'next/link' -import { PropsWithChildren, useState } from 'react' +import { PropsWithChildren } from 'react' import { TooltipProvider } from '~/components/ui/tooltip' import { useDatabasesQuery } from '~/data/databases/databases-query' -import { useBreakpoint } from '~/lib/use-breakpoint' import { currentDomainHostname, currentDomainUrl, @@ -18,7 +16,6 @@ import { legacyDomainUrl, } from '~/lib/util' import { useApp } from './app-provider' -import { LiveShareIcon } from './live-share-icon' import Sidebar from './sidebar' import { Accordion, AccordionContent, AccordionItem, AccordionTrigger } from './ui/accordion' import { Dialog, DialogContent, DialogHeader, DialogTitle } from './ui/dialog' @@ -29,14 +26,12 @@ export type LayoutProps = PropsWithChildren export default function Layout({ children }: LayoutProps) { const { isLegacyDomain, isLegacyDomainRedirect } = useApp() - const isSmallBreakpoint = useBreakpoint('lg') return (

- {!isLegacyDomain && } {(isLegacyDomain || isLegacyDomainRedirect) && }
@@ -52,86 +47,6 @@ export default function Layout({ children }: LayoutProps) { ) } -function LiveShareBanner() { - const [videoLoaded, setVideoLoaded] = useState(false) - - return ( -
- New: Connect to your in-browser databases from outside the browser. - setVideoLoaded(false)}> - - Learn more. - - - - Introducing Live Share -
- - -
-

- With Live Share, you can connect directly to your in-browser PGlite databases from{' '} - outside the browser. -

-
- {!videoLoaded && ( -
- -
- )} - setVideoLoaded(true)} - > - - -
- - -

How does it work?

- -
    -
  1. - Click on the {' '} - menu next your database and tap{' '} - - Live Share - -
    2. -
  2. A unique connection string will appear for your database
    3. -
  3. - Copy-paste the connection string into any Postgres client (like psql) - and begin querying! -
    4. -
- -
- -
-
- ) -} - function RenameBanner() { const { setIsRenameDialogOpen } = useApp() return ( diff --git a/apps/postgres-new/components/model-provider/set-model-provider-dialog.tsx b/apps/postgres-new/components/model-provider/set-model-provider-dialog.tsx new file mode 100644 index 00000000..0ca8758a --- /dev/null +++ b/apps/postgres-new/components/model-provider/set-model-provider-dialog.tsx @@ -0,0 +1,266 @@ +'use client' + +import { zodResolver } from '@hookform/resolvers/zod' +import { m } from 'framer-motion' +import { Brain, Expand } from 'lucide-react' +import { useState } from 'react' +import { useForm, useWatch } from 'react-hook-form' +import { z } from 'zod' +import { useApp } from '~/components/app-provider' +import { Button, ButtonProps } from '~/components/ui/button' +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from '~/components/ui/dialog' +import { + Form, + FormControl, + FormField, + FormItem, + FormLabel, + FormMessage, +} from '~/components/ui/form' +import { Input } from '~/components/ui/input' +import { Switch } from '~/components/ui/switch' +import { Textarea } from '~/components/ui/textarea' +import { getProviderUrl } from '~/lib/llm-provider' +import { getSystemPrompt } from '~/lib/system-prompt' + +const formSchema = z.object({ + apiKey: z + .string() + .transform((str) => (str === '' ? undefined : str)) + .optional(), + baseUrl: z.string().min(1), + model: z.string().min(1), + system: z.string().min(1), + enabled: z.boolean(), +}) + +type FormSchema = z.infer + +function SetModelProviderForm(props: { id: string; onSubmit: (values: FormSchema) => void }) { + const { modelProvider } = useApp() + + const form = useForm({ + resolver: zodResolver(formSchema), + defaultValues: { + enabled: false, + system: getSystemPrompt(), + ...modelProvider.state, + }, + }) + + const isEnabled = useWatch({ control: form.control, name: 'enabled' }) + + const [isPromptExpanded, setIsPromptExpanded] = useState(false) + + async function onSubmit(values: z.infer) { + await modelProvider.set(values) + props.onSubmit(values) + } + + return ( +
+ + ( + +
+ + + + Enable +
+ + + )} + /> + {isEnabled && ( + <> + ( + + Base URL + + <> + +
+ { + e.preventDefault() + form.setValue('baseUrl', getProviderUrl('openai')) + form.setValue('model', 'gpt-4o') + }} + > + OpenAI + + { + e.preventDefault() + form.setValue('baseUrl', getProviderUrl('x-ai')) + form.setValue('model', 'grok-beta') + }} + > + xAI + + { + e.preventDefault() + field.onChange({ target: { value: getProviderUrl('openrouter') } }) + form.setValue('model', '') + }} + > + OpenRouter + +
+ + + + + )} + /> + ( + + API key + + + + + + )} + /> + ( + + Model + + + + + + )} + /> + ( + + System prompt + + <> + { + setIsPromptExpanded(true) + }} + > +
+ {field.value} +
+
+ + + {isPromptExpanded && ( +
+ +